Middleware Chain

middlewares = [
    ThreadDataMiddleware(),
    UploadsMiddleware(),
    SandboxMiddleware(),
    DanglingToolCallMiddleware(),
    # Conditionally added:
    SummarizationMiddleware(),  # if enabled
    TodoListMiddleware(),        # if is_plan_mode
    TitleMiddleware(),
    MemoryMiddleware(),
    ViewImageMiddleware(),       # if model supports vision
    SubagentLimitMiddleware(),   # if subagent_enabled
    ClarificationMiddleware()    # must be last
]

class ThreadDataMiddleware(AgentMiddleware[ThreadDataMiddlewareState]):
    def __init__(self, base_dir: str | None = None, lazy_init: bool = True):
        # lazy_init=True: Only compute paths, defer directory creation
        # lazy_init=False: Eagerly create directories
        self._paths = Paths(base_dir) if base_dir else get_paths()
        self._lazy_init = lazy_init

    def before_agent(self, state, runtime):
        thread_id = runtime.context.get("thread_id")
        if self._lazy_init:
            paths = self._get_thread_paths(thread_id)
        else:
            paths = self._create_thread_directories(thread_id)
        
        return {
            "thread_data": {
                "workspace_path": str(paths["workspace_path"]),
                "uploads_path": str(paths["uploads_path"]),
                "outputs_path": str(paths["outputs_path"])
            }
        }

backend/.deer-flow/threads/{thread_id}/user-data/
├── workspace/  # Agent's working directory
├── uploads/    # User-uploaded files
└── outputs/    # Files presented to user via present_files tool

class UploadsMiddleware(AgentMiddleware[UploadsMiddlewareState]):
    def before_agent(self, state, runtime):
        thread_id = runtime.context.get("thread_id")
        messages = state.get("messages", [])
        
        # Track previously shown files from message history
        shown_files = self._extract_files_from_previous_messages(messages[:-1])
        
        # List only newly uploaded files
        new_files = self._list_newly_uploaded_files(thread_id, shown_files)
        
        if new_files:
            # Prepend file list to last human message
            files_message = self._create_files_message(new_files)
            updated_message = HumanMessage(
                content=f"{files_message}\n\n{original_content}"
            )
            messages[-1] = updated_message
        
        return {"uploaded_files": new_files, "messages": messages}

class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
    def __init__(self, lazy_init: bool = True):
        # lazy_init=True: Acquire on first tool call
        # lazy_init=False: Acquire in before_agent()
        self._lazy_init = lazy_init

    def before_agent(self, state, runtime):
        if self._lazy_init:
            return None  # Defer acquisition
        
        if "sandbox" not in state or state["sandbox"] is None:
            thread_id = runtime.context["thread_id"]
            sandbox_id = self._acquire_sandbox(thread_id)
            return {"sandbox": {"sandbox_id": sandbox_id}}

class DanglingToolCallMiddleware(AgentMiddleware[AgentState]):
    def wrap_model_call(self, request, handler):
        # Scan for AIMessages with tool_calls that lack ToolMessage responses
        patched = self._build_patched_messages(request.messages)
        if patched:
            request = request.override(messages=patched)
        return handler(request)
    
    def _build_patched_messages(self, messages):
        existing_tool_msg_ids = {msg.tool_call_id for msg in messages 
                                  if isinstance(msg, ToolMessage)}
        
        patched = []
        for msg in messages:
            patched.append(msg)
            if getattr(msg, "type", None) == "ai":
                for tc in getattr(msg, "tool_calls", []):
                    if tc["id"] not in existing_tool_msg_ids:
                        # Inject placeholder ToolMessage
                        patched.append(ToolMessage(
                            content="[Tool call was interrupted and did not return a result.]",
                            tool_call_id=tc["id"],
                            status="error"
                        ))
        return patched

class SummarizationConfig(BaseModel):
    enabled: bool = False
    model_name: str | None = None  # None = use lightweight model
    trigger: ContextSize | list[ContextSize] | None
    keep: ContextSize = ContextSize(type="messages", value=20)
    trim_tokens_to_summarize: int | None = 4000
    summary_prompt: str | None = None

def _create_summarization_middleware():
    config = get_summarization_config()
    if not config.enabled:
        return None
    
    # Convert config to middleware parameters
    trigger = [t.to_tuple() for t in config.trigger] if isinstance(config.trigger, list) 
              else config.trigger.to_tuple()
    keep = config.keep.to_tuple()
    model = config.model_name or create_chat_model(thinking_enabled=False)
    
    return SummarizationMiddleware(
        model=model,
        trigger=trigger,
        keep=keep,
        trim_tokens_to_summarize=config.trim_tokens_to_summarize
    )

def _create_todo_list_middleware(is_plan_mode: bool):
    if not is_plan_mode:
        return None
    
    system_prompt = """
    <todo_list_system>
    **CRITICAL RULES:**
    - Mark todos as completed IMMEDIATELY after finishing each step
    - Keep EXACTLY ONE task as `in_progress` at any time
    - Update in REAL-TIME - gives users visibility
    - DO NOT use for simple tasks (< 3 steps)
    </todo_list_system>
    """
    
    tool_description = """Use for complex tasks (3+ steps) only..."""
    
    return TodoListMiddleware(
        system_prompt=system_prompt,
        tool_description=tool_description
    )

class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
    def after_agent(self, state, runtime):
        if self._should_generate_title(state):
            title = self._generate_title(state)
            return {"title": title}
        return None
    
    def _should_generate_title(self, state):
        config = get_title_config()
        if not config.enabled or state.get("title"):
            return False
        
        messages = state.get("messages", [])
        user_messages = [m for m in messages if m.type == "human"]
        assistant_messages = [m for m in messages if m.type == "ai"]
        
        # Generate after first complete exchange
        return len(user_messages) == 1 and len(assistant_messages) >= 1
    
    def _generate_title(self, state):
        config = get_title_config()
        model = create_chat_model(thinking_enabled=False)  # Lightweight model
        
        user_msg = next(m.content for m in messages if m.type == "human")
        assistant_msg = next(m.content for m in messages if m.type == "ai")
        
        prompt = config.prompt_template.format(
            max_words=config.max_words,
            user_msg=user_msg[:500],
            assistant_msg=assistant_msg[:500]
        )
        
        response = model.invoke(prompt)
        title = response.content.strip()[:config.max_chars]
        return title

class MemoryMiddleware(AgentMiddleware[MemoryMiddlewareState]):
    def __init__(self, agent_name: str | None = None):
        # agent_name: If provided, uses per-agent memory storage
        self._agent_name = agent_name
    
    def after_agent(self, state, runtime):
        config = get_memory_config()
        if not config.enabled:
            return None
        
        thread_id = runtime.context.get("thread_id")
        messages = state.get("messages", [])
        
        # Filter to user inputs + final assistant responses (no tool calls)
        filtered_messages = _filter_messages_for_memory(messages)
        
        # Queue for debounced background processing
        queue = get_memory_queue()
        queue.add(
            thread_id=thread_id,
            messages=filtered_messages,
            agent_name=self._agent_name
        )
        
        return None  # No state changes

def _filter_messages_for_memory(messages):
    filtered = []
    for msg in messages:
        if msg.type == "human":
            filtered.append(msg)  # Always keep user messages
        elif msg.type == "ai" and not getattr(msg, "tool_calls", None):
            filtered.append(msg)  # Only keep final AI responses
    return filtered

class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
    def before_model(self, state, runtime):
        return self._inject_image_message(state)
    
    def _should_inject_image_message(self, state):
        messages = state.get("messages", [])
        last_assistant_msg = self._get_last_assistant_message(messages)
        
        if not last_assistant_msg:
            return False
        
        # Check if it has view_image tool calls
        if not self._has_view_image_tool(last_assistant_msg):
            return False
        
        # Check if all tools completed
        if not self._all_tools_completed(messages, last_assistant_msg):
            return False
        
        # Check if we already injected the message
        # (prevents duplicate injections)
        return not self._already_injected(messages, last_assistant_msg)
    
    def _create_image_details_message(self, state):
        viewed_images = state.get("viewed_images", {})
        content_blocks = [
            {"type": "text", "text": "Here are the images you've viewed:"}
        ]
        
        for image_path, image_data in viewed_images.items():
            content_blocks.append({
                "type": "text",
                "text": f"\n- **{image_path}** ({image_data['mime_type']})"
            })
            content_blocks.append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:{image_data['mime_type']};base64,{image_data['base64']}"
                }
            })
        
        return content_blocks

def merge_viewed_images(existing, new):
    if new == {}:  # Empty dict clears all viewed images
        return {}
    return {**existing, **new}  # Merge dictionaries

class SubagentLimitMiddleware(AgentMiddleware[AgentState]):
    def __init__(self, max_concurrent: int = MAX_CONCURRENT_SUBAGENTS):
        # max_concurrent clamped to [2, 4]
        self.max_concurrent = _clamp_subagent_limit(max_concurrent)
    
    def after_model(self, state, runtime):
        return self._truncate_task_calls(state)
    
    def _truncate_task_calls(self, state):
        messages = state.get("messages", [])
        last_msg = messages[-1]
        
        if getattr(last_msg, "type", None) != "ai":
            return None
        
        tool_calls = getattr(last_msg, "tool_calls", None)
        task_indices = [i for i, tc in enumerate(tool_calls) 
                        if tc.get("name") == "task"]
        
        if len(task_indices) <= self.max_concurrent:
            return None
        
        # Keep only first max_concurrent task calls
        indices_to_drop = set(task_indices[self.max_concurrent:])
        truncated = [tc for i, tc in enumerate(tool_calls) 
                     if i not in indices_to_drop]
        
        logger.warning(f"Truncated {len(indices_to_drop)} excess task calls")
        
        updated_msg = last_msg.model_copy(update={"tool_calls": truncated})
        return {"messages": [updated_msg]}

class ClarificationMiddleware(AgentMiddleware[ClarificationMiddlewareState]):
    def wrap_tool_call(self, request, handler):
        if request.tool_call.get("name") != "ask_clarification":
            return handler(request)  # Pass through
        
        return self._handle_clarification(request)
    
    def _handle_clarification(self, request):
        args = request.tool_call.get("args", {})
        formatted_message = self._format_clarification_message(args)
        
        tool_message = ToolMessage(
            content=formatted_message,
            tool_call_id=request.tool_call.get("id"),
            name="ask_clarification"
        )
        
        # Return Command that interrupts execution
        return Command(
            update={"messages": [tool_message]},
            goto=END  # Stop execution, wait for user response
        )
    
    def _format_clarification_message(self, args):
        question = args.get("question", "")
        clarification_type = args.get("clarification_type", "missing_info")
        context = args.get("context")
        options = args.get("options", [])
        
        type_icons = {
            "missing_info": "❓",
            "ambiguous_requirement": "🤔",
            "approach_choice": "🔀",
            "risk_confirmation": "⚠️",
            "suggestion": "💡"
        }
        
        icon = type_icons.get(clarification_type, "❓")
        
        message_parts = []
        if context:
            message_parts.append(f"{icon} {context}")
            message_parts.append(f"\n{question}")
        else:
            message_parts.append(f"{icon} {question}")
        
        if options:
            message_parts.append("")
            for i, option in enumerate(options, 1):
                message_parts.append(f"  {i}. {option}")
        
        return "\n".join(message_parts)

config = {
    "configurable": {
        "thinking_enabled": True,
        "model_name": "gpt-4o",
        "is_plan_mode": False,      # Enables TodoListMiddleware
        "subagent_enabled": True,    # Enables SubagentLimitMiddleware
        "max_concurrent_subagents": 3
    }
}

agent = make_lead_agent(config)

class ThreadState(AgentState):
    sandbox: NotRequired[SandboxState | None]
    thread_data: NotRequired[ThreadDataState | None]
    title: NotRequired[str | None]
    artifacts: Annotated[list[str], merge_artifacts]
    todos: NotRequired[list | None]
    uploaded_files: NotRequired[list[dict] | None]
    viewed_images: Annotated[dict[str, ViewedImageData], merge_viewed_images]

logger.warning(f"Injecting {count} placeholder ToolMessage(s) for dangling tool calls")
logger.warning(f"Truncated {count} excess task call(s) from model response (limit: {limit})")
print(f"[ViewImageMiddleware] Injecting image details message with images before LLM call")
print(f"[ClarificationMiddleware] Intercepted clarification request")

cd backend
make dev  # Watch logs in terminal