vamp-cli 0.1.0__py3-none-any.whl

vamp_cli/agent/context_sync.py

@@ -0,0 +1,72 @@
+from langchain_core.messages import SystemMessage, BaseMessage
+
+def mode_switch_sync(
+    prev_mode: str,
+    new_mode: str,
+    last_state = None,
+) -> dict:
+    """Return payload additions to inject when switching modes.
+    Returns {} if there is nothing to sync.
+    """
+    if prev_mode == new_mode or last_state is None:
+        return {}
+
+    # plan -> normal : hand the plan to the coder as context
+    if prev_mode == "plan" and new_mode == "normal":
+        plan = last_state.get("plan") or "(no plan was produced)"
+        # todos = last_state.get("todo") or []
+        # todo_lines = "\n".join(
+        #     f"  [{t['status']}] {t['id']}. {t['description']}" for t in todos
+        # ) or "  (empty)"
+
+        sync_msg = SystemMessage(
+            content=(
+                "MODE SWITCH: The user has switched you from PLAN mode to NORMAL mode.\n"
+                "You are no longer bound to a planner. You are free to act directly.\n"
+                "For your reference, this is the plan that was produced earlier:\n\n"
+                f"{plan}\n\n"
+                # "Current todo list at the moment of the switch:\n"
+                # f"{todo_lines}\n\n"
+                "Treat this as context only; do not re-inject the plan verbatim."
+                "NOTE: IF YOU HAVE ALREADY EXECUTED THIS PLAN THEN DO NOT CONSIDER THIS."
+            )
+        )
+        return {"coder_messages": [sync_msg]}
+
+    # normal -> plan : summarise coder activity for the planner
+    if prev_mode == "normal" and new_mode == "plan":
+        coder_msgs: list[BaseMessage] = last_state.get("coder_messages", [])
+        summary = _summarise_coder_activity(coder_msgs)
+
+        sync_msg = SystemMessage(
+            content=(
+                "MODE SWITCH: The user has switched you from NORMAL mode to PLAN mode.\n"
+                "Below is a compact summary of what the coding agent already did in "
+                "the previous (normal) mode. Build your plan on top of this — do not "
+                "repeat finished work.\n\n"
+                f"{summary}"
+            )
+        )
+        return {"planner_messages": [sync_msg]}
+
+    return {}
+
+
+def _summarise_coder_activity(msgs: list[BaseMessage]) -> str:
+    """
+    Cheap deterministic summary: last N AI text contents + the names of any
+    tools called.
+    """
+    lines = []
+    for m in msgs[-10:]:                     # last N turns only
+        # AI text
+        content = getattr(m, "content", "")
+        if content:
+            text = content if isinstance(content, str) else str(content)
+            lines.append(f"AI: {text[:400]}")
+
+        # tool calls
+        for tc in (getattr(m, "tool_calls", None) or []):
+            lines.append(f"  -> tool: {tc['name']}")
+
+    return "\n".join(lines) if lines else "(no prior activity recorded)"

vamp_cli/agent/graph.py

@@ -0,0 +1,159 @@
+from typing import Literal
+
+from langgraph.graph import StateGraph, START, END
+from langchain_core.messages import AIMessage
+from langgraph.checkpoint.memory import MemorySaver
+
+from vamp_cli.tools.registry import ToolRegistry
+from vamp_cli.agent.state import AgentState
+from vamp_cli.agent.nodes import (
+    create_agent_node,
+    create_planner_node,
+    create_plan_approval_node,
+    create_tool_classify_node,
+    create_tool_approve_node,
+    create_tool_execute_node,
+)
+
+MAX_ITERATIONS = 15
+
+
+def route_start(state) -> Literal["agent", "planner"]:
+    return "planner" if state["mode"] == "plan" else "agent"
+
+
+def should_continue(state) -> Literal["classify", "__end__"]:
+    if state["iteration"] >= MAX_ITERATIONS:
+        return "__end__"
+    last = state["coder_messages"][-1]
+    if not isinstance(last, AIMessage) or not last.tool_calls:
+        return "__end__"
+    return "classify"
+
+
+def route_after_planner(state) -> Literal["classify", "__end__"]:
+    if state["iteration"] >= MAX_ITERATIONS:
+        return "__end__"
+    msgs = state["planner_messages"]
+    last = next((m for m in reversed(msgs) if isinstance(m, AIMessage)), None)
+    if last is None or not last.tool_calls:
+        return "__end__"
+    return "classify"
+
+
+def route_after_execute(state) -> Literal["coder", "planner", "plan_approver"]:
+    """Reads the (still-populated) pipeline channel + last planner AIMessage."""
+    channel = state.get("_pending_tool_channel")
+
+    if channel == "planner_messages":
+        msgs = state["planner_messages"]
+        last_ai = next((m for m in reversed(msgs) if isinstance(m, AIMessage)), None)
+        if last_ai and last_ai.tool_calls:
+            if any(tc["name"] == "submit_plan" for tc in last_ai.tool_calls):
+                return "plan_approver"
+        return "planner"
+
+    return "coder"
+
+
+def route_after_plan_approver(state) -> Literal["coder", "__end__"]:
+    return "coder" if state.get("requires_execution") else "__end__"
+
+
+def create_graph(
+    coder_model, 
+    planner_model, 
+    tool_registry: ToolRegistry,
+    checkpointer=None,
+):
+
+    coder_node    = create_agent_node(coder_model)
+    planner_node  = create_planner_node(planner_model)
+    approver_node = create_plan_approval_node(planner_model)
+
+    coder_classify   = create_tool_classify_node("coder_messages")
+    planner_classify = create_tool_classify_node("planner_messages")
+    
+    approve_tools    = create_tool_approve_node()
+    execute_tools    = create_tool_execute_node(tool_registry)
+
+    g = StateGraph(AgentState)
+
+    g.add_node("coder", coder_node)
+    g.add_node("planner", planner_node)
+    
+    g.add_node("plan_approver", approver_node)
+    
+    g.add_node("coder_classify", coder_classify)
+    g.add_node("planner_classify", planner_classify)
+    
+    g.add_node("approve_tools", approve_tools)
+    g.add_node("execute_tools", execute_tools)
+
+    # START
+    g.add_conditional_edges(
+        START, 
+        route_start,
+        {
+            "agent": "coder", 
+            "planner": "planner"
+        }
+    )
+
+    # PLANNER -> classify -> approve -> execute -> route
+    g.add_conditional_edges(
+        "planner", 
+        route_after_planner,
+        {
+            "classify": "planner_classify", 
+            "__end__": END
+        }
+    )
+    
+    g.add_edge(
+        "planner_classify", 
+        "approve_tools"
+    )
+
+    # CODER -> classify -> approve -> execute -> route
+    g.add_conditional_edges(
+        "coder", 
+        should_continue,
+        {
+            "classify": "coder_classify", 
+            "__end__": END
+        }
+    )
+    
+    g.add_edge(
+        "coder_classify", 
+        "approve_tools"
+    )
+
+    # shared tail
+    g.add_edge(
+        "approve_tools", 
+        "execute_tools"
+    )
+    
+    g.add_conditional_edges(
+        "execute_tools", 
+        route_after_execute, 
+        {
+            "coder": "coder",
+            "planner": "planner",
+            "plan_approver": "plan_approver",
+        }
+    )
+
+    # plan approval gate
+    g.add_conditional_edges(
+        "plan_approver", 
+        route_after_plan_approver,
+        {
+            "coder": "coder", 
+            "__end__": END
+        }
+    )
+
+    return g.compile(checkpointer=checkpointer or MemorySaver())

vamp_cli/agent/instructions.py

@@ -0,0 +1,18 @@
+TODO_INSTRUCTIONS = """\
+    You maintain a todo list for every non-trivial task, using the `write_update_todos` tool.
+
+    Rules:
+    - Call `write_update_todos` BEFORE you start making changes, and again whenever the
+    status of any item changes.
+    - Always send the COMPLETE list — it replaces the previous one.
+    - Exactly one item may be `in_progress` at a time.
+    - Mark an item `completed` as soon as it is verified done, then set the next
+    item to `in_progress` BEFORE starting it.
+    - Never mark something `completed` if verification failed or it is only
+    partially done — use `blocked` instead.
+    - Keep item ids stable across updates.
+"""
+
+PLAN_INSTRUCTIONS = """\
+    The plan above was produced by a separate planning agent after repository exploration. Treat it as the intended implementation plan, not as the user's original request. So execute the plan
+"""

vamp_cli/agent/nodes.py

@@ -0,0 +1,375 @@
+import json
+import time
+import asyncio
+from typing import Any
+from prompt_toolkit import print_formatted_text
+from prompt_toolkit.formatted_text import FormattedText, HTML
+
+from langgraph.graph import END
+from langgraph.types import interrupt, Command
+from langchain_core.runnables import Runnable, RunnableConfig
+from langchain_core.messages import ToolMessage, SystemMessage, AIMessage, BaseMessage
+
+from vamp_cli.tools.registry import ToolRegistry
+from vamp_cli.agent.state import AgentState, TodoItem
+from vamp_cli.agent.state import AgentState
+from vamp_cli.agent.status import FeedbackResponse
+from vamp_cli.agent.instructions import TODO_INSTRUCTIONS, PLAN_INSTRUCTIONS
+
+
+
+def create_planner_node(model: Runnable):
+
+    async def planner_node(state: AgentState, config: RunnableConfig):
+        # The planner gets the same model or different model (if user as configured for planner explicitly) interface,
+        # with a restricted read-only tool registry.
+
+        messages = state["planner_messages"]
+
+        # response = await model.ainvoke(messages)
+        response = await model.with_config(config).ainvoke(messages)
+        pending = next(
+            (tc for tc in (response.tool_calls or []) if tc["name"] == "submit_plan"), 
+            None
+        )
+        
+        if response.tool_calls:
+            print_formatted_text(
+                FormattedText([(
+                    "dim",
+                    "\n\n[planner] calling:"
+                )]),
+                end="\n"
+            )
+            for tc in response.tool_calls:
+                print_formatted_text(
+                    FormattedText([(
+                        "dim", 
+                        f">> {tc['name']}({tc['args']})"
+                    )]),
+                    end="\n"
+                )
+                # print("Tool: ", tc['name'])
+                # print("Arguments: ", tc['args'])
+        
+        return {
+            "planner_messages": [response],
+            "plan": pending["args"].get("plan") if pending else None,
+            "iteration": state["iteration"] + 1,
+        }
+    
+    return planner_node
+
+
+def create_plan_approval_node(model: Runnable):
+    
+    async def plan_approval_node(state: AgentState, config: RunnableConfig):
+        
+        if not state.get("plan"):
+            return {"requires_execution": False}
+        
+        checkpoint_payload = {
+            "implementation_permission": True,
+            "plan": state["plan"],
+            "warning": "Asking before implementing this plan"
+        }
+                    
+        # execution halts here
+        fb: FeedbackResponse = interrupt(value=checkpoint_payload, response_schema=FeedbackResponse)
+        return {"requires_execution": fb.approved}   
+    
+    return plan_approval_node
+
+
+
+def create_agent_node(model: Runnable):
+
+    async def agent_node(state: AgentState, config: RunnableConfig):
+        
+        message_to_send = list(state["coder_messages"])
+        
+        if state["mode"] == "plan":
+            plan = state["plan"]
+            
+            if plan:
+                plan_message = SystemMessage(content=f"Here is the project plan:\n{state['plan']}")
+                plan_instruction = SystemMessage(content=PLAN_INSTRUCTIONS)
+                
+                message_to_send += [plan_message] + [plan_instruction]
+                
+                
+                
+        if not state.get("todo"):
+            message_to_send.append(SystemMessage(content=TODO_INSTRUCTIONS))
+            
+        if state.get("todo"):
+            turns_since_update = _turns_since_todo_update(state["coder_messages"])
+            if turns_since_update >= 10:
+                message_to_send.append(
+                    SystemMessage(content=_format_todo_reminder(state["todo"]))
+                )
+        
+        response: AIMessage = await model.with_config(config).ainvoke(message_to_send)
+        
+        # print("\n\nCODER------AI-RESPONSE:\n.\n.\n", response, "\n\n")
+        
+        if response.tool_calls:
+            print_formatted_text(
+                FormattedText([(
+                    "dim",
+                    "\n\n[coder] calling:"
+                )]),
+                end="\n"
+            )
+            for tc in response.tool_calls:
+                print_formatted_text(
+                    FormattedText([(
+                        "dim", 
+                        f">> {tc['name']}({tc['args']})"
+                    )]),
+                    end="\n"
+                )
+        
+        return {
+            "coder_messages": [response],
+            # "is_plan_injected": plan_injected,
+            "iteration": state["iteration"] + 1,
+        }
+    return agent_node
+
+
+DANGEROUS_PATTERNS = ("rm ", "rmdir ", "curl ", "wget ", "pip install ", "shutdown", "dd ")
+
+def _is_dangerous(tool_name: str, tool_args: dict) -> bool:
+    is_bash = tool_name == "execute_command" or "execute" in tool_name.lower()
+    if not is_bash:
+        return False
+    cmd = str(tool_args.get("command", "")).lower()
+    return any(p in cmd for p in DANGEROUS_PATTERNS)
+
+
+def create_tool_classify_node(channel: str):
+    async def tool_classify_node(state: AgentState, config: RunnableConfig):
+        last = state[channel][-1]
+
+        if not isinstance(last, AIMessage) or not last.tool_calls:
+            raise RuntimeError(
+                f"classify_node invoked but last {channel} entry "
+                f"is {type(last).__name__} with no tool calls"
+            )
+
+        dangerous_ids = [
+            tc["id"]
+            for tc in last.tool_calls
+            if _is_dangerous(tc["name"], tc.get("args") or {})
+        ]
+
+        return {
+            "_pending_tool_channel":    channel,
+            "_pending_tool_calls":      [dict(tc) for tc in last.tool_calls],
+            "_dangerous_tool_call_ids": dangerous_ids,
+            "_approved_tool_call_ids":  [],
+            "_denied_tool_call_ids":    [],
+        }
+
+    return tool_classify_node
+
+
+def create_tool_approve_node():
+    async def tool_approve_node(state: AgentState, config: RunnableConfig):
+        pending       = state.get("_pending_tool_calls", [])
+        dangerous_ids = state.get("_dangerous_tool_call_ids", [])
+        all_ids       = [tc["id"] for tc in pending]
+
+        # No dangerous calls: auto-approve everything, no interrupt.
+        if not dangerous_ids:
+            return {
+                "_approved_tool_call_ids": all_ids,
+                "_denied_tool_call_ids":   [],
+            }
+
+        dangerous_calls = [tc for tc in pending if tc["id"] in dangerous_ids]
+
+        payload = {
+            "requires_tool_approval": True,
+            "dangerous_calls": [
+                {
+                    "tool_call_id": tc["id"],
+                    "tool_name":    tc["name"],
+                    "tool_args":    tc.get("args") or {},
+                }
+                for tc in dangerous_calls
+            ],
+            "safe_call_count": len(pending) - len(dangerous_calls),
+            "warning": (
+                "The following tool calls were flagged as dangerous. "
+                "Approving will execute them (plus all safe calls) exactly once."
+            ),
+        }
+
+        # Halts here. On resume, returns the FeedbackResponse instead of raising.
+        fb: FeedbackResponse = interrupt(value=payload, response_schema=FeedbackResponse)
+
+        safe_ids = [i for i in all_ids if i not in dangerous_ids]
+
+        if fb.approved:
+            return {
+                "_approved_tool_call_ids": safe_ids + dangerous_ids,
+                "_denied_tool_call_ids":   [],
+            }
+        return {
+            "_approved_tool_call_ids": safe_ids,
+            "_denied_tool_call_ids":   dangerous_ids,
+        }
+
+    return tool_approve_node
+
+
+def create_tool_execute_node(tool_registry: ToolRegistry):
+    async def tool_execute_node(state: AgentState, config: RunnableConfig):
+        channel = state["_pending_tool_channel"]
+        pending = state.get("_pending_tool_calls", [])
+        approved = set(state.get("_approved_tool_call_ids", []))
+        denied = set(state.get("_denied_tool_call_ids", []))
+
+        tool_messages: list[ToolMessage] = []
+        todo_update = None
+        
+        interrupted = False
+        
+        emitter = (config.get("configurable") or {}).get("emitter")
+        node_name = channel.replace("_messages", "") if channel else None
+
+        for tc in pending:
+            tc_id, tc_name, tc_args = tc["id"], tc["name"], tc.get("args") or {}
+
+            if tc_id in denied:
+                content=(
+                    "CRITICAL ACCESS DENIED: The security administrator has "
+                    "EXPLICITLY FORBIDDEN executing the denied task in any way. "
+                    "MANDATE: cease all attempts to execute the denied task. "
+                    "Abort this task immediately and explain the restriction "
+                    "to the user. Do not retry. Move on to another task."
+                )
+                
+                tool_messages.append(ToolMessage(
+                    content=content,
+                    tool_call_id=tc_id,
+                    name=tc_name,
+                    status="error",
+                ))
+                
+                if emitter is not None:
+                    emitter.emit_tool_result(
+                        tool_call_id=tc_id, tool_name=tc_name, node=node_name,
+                        status="denied", content=content, duration_ms=0,
+                    )
+                continue
+
+            if tc_id not in approved:
+                content = f"Tool '{tc_name}' was not approved for execution."
+                tool_messages.append(ToolMessage(
+                    content=content,
+                    tool_call_id=tc_id, name=tc_name, status="error",
+                ))
+                if emitter is not None:
+                    emitter.emit_tool_result(
+                        tool_call_id=tc_id, tool_name=tc_name, node=node_name,
+                        status="error", content=content, duration_ms=0,
+                    )
+                continue
+            
+            if interrupted:
+                content="[Skipped — session was interrupted by user before this tool ran]"
+                tool_messages.append(ToolMessage(
+                    content=content,
+                    tool_call_id=tc_id, name=tc_name, status="error",
+                ))
+                
+                if emitter is not None:
+                    emitter.emit_tool_result(
+                        tool_call_id=tc_id, tool_name=tc_name, node=node_name,
+                        status="error", content=content, duration_ms=0,
+                    )
+                continue
+            
+            start = time.time()
+            try:
+                result = await tool_registry.call(tc_name, tc_args)
+            except (asyncio.CancelledError, KeyboardInterrupt):
+                interrupted = True
+                result = "[Tool execution was interrupted by user]"
+                status = "error"
+            except Exception as e:
+                result = {"error": f"Tool '{tc_name}' failed: {e}"}
+                status = "error"
+            duration_ms = int((time.time() - start) * 1000)
+
+            if tc_name == "write_update_todos" and isinstance(result, str) and not interrupted:
+                try:
+                    todo_update = json.loads(result)
+                except json.JSONDecodeError:
+                    todo_update = None
+
+            string_result = result if isinstance(result, str) else json.dumps(result)
+            tool_messages.append(ToolMessage(
+                content=string_result,
+                tool_call_id=tc_id,
+                name=tc_name,
+                status="success" if isinstance(result, str) else "error",
+            ))
+            if emitter is not None:
+                emitter.emit_tool_result(
+                    tool_call_id=tc_id, tool_name=tc_name, node=node_name,
+                    status="success", content=string_result,
+                    duration_ms=duration_ms,
+                )
+
+        update = {
+            channel: tool_messages,
+            "_pending_tool_calls":     [],
+            "_dangerous_tool_call_ids": [],
+            "_approved_tool_call_ids":  [],
+            "_denied_tool_call_ids":    [],
+        }
+        if todo_update is not None:
+            update["todo"] = todo_update
+            
+        # if interrupted:
+        #     # Save the state update, then END the graph invocation.
+        #     # Control returns to main(); the prompt shows again.
+        #     return Command(update=update, goto=END)
+        
+        return update
+
+    return tool_execute_node
+
+
+def _turns_since_todo_update(messages: list[BaseMessage]) -> int:
+    """How many AIMessages have passed since the last write_update_todos call."""
+    count = 0
+    for msg in reversed(messages):
+        if isinstance(msg, AIMessage) and msg.tool_calls:
+            if any(tc["name"] == "write_update_todos" for tc in msg.tool_calls):
+                return count
+            count += 1
+        elif isinstance(msg, AIMessage):
+            count += 1
+    return count
+
+
+def _format_todo_reminder(todos: list[TodoItem]) -> str:
+    marks = {
+        "pending": "[ ]",
+        "in_progress": "[~]",
+        "completed": "[x]",
+        "blocked": "[!]",
+    }
+    lines = ["Your current todo list:"]
+    for t in todos:
+        lines.append(f"  {marks[t['status']]} {t['id']}. {t['description']}")
+    lines.append(
+        "Continue with the next in_progress or pending item. "
+        "Call write_update_todos when the status of any item changes."
+    )
+    return "\n".join(lines)

vamp_cli/agent/roles.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+class AgentRole(Enum):
+    DEFAULT = "default"
+    DEFAULT_PLANNER = "default_planner"
+    PLANNER = "planner"
+    CODER = "coder"
+