Agent: : Update ToolRunContext and Execution Hooks for PlanV2 Nativity

portia/__init__.py

@@ -91,7 +91,7 @@
 
 # Plan and execution related classes
 from portia.plan import Plan, PlanBuilder, PlanContext, PlanInput, PlanUUID, Step, Variable
-from portia.plan_run import PlanRun, PlanRunState
+from portia.plan_run import PlanRun, PlanRunState, PlanRunV2, StepOutputValue
 
 # Core classes
 from portia.portia import ExecutionHooks, Portia
@@ -163,6 +163,7 @@
     "PlanRun",
     "PlanRunNotFoundError",
     "PlanRunState",
+    "PlanRunV2",
     "PlanUUID",
     "PlanV2",
     "PlanningAgentType",
@@ -176,6 +177,7 @@
     "StdioMcpClientConfig",
     "Step",
     "StepOutput",
+    "StepOutputValue",
     "StepV2",
     "StorageClass",
     "StorageError",

portia/execution_hooks.py

@@ -8,6 +8,8 @@
 
 from pydantic import BaseModel, ConfigDict
 
+from portia.builder.plan_v2 import PlanV2
+from portia.builder.step_v2 import StepV2
 from portia.clarification import (
     Clarification,
     ClarificationCategory,
@@ -18,8 +20,7 @@
 from portia.errors import ToolHardError
 from portia.execution_agents.output import Output
 from portia.logger import logger
-from portia.plan import Plan, Step
-from portia.plan_run import PlanRun
+from portia.plan_run import PlanRunV2
 from portia.tool import Tool
 
 
@@ -53,75 +54,75 @@
     clarification_handler: ClarificationHandler | None = None
     """Handler for clarifications raised during execution."""
 
-    before_step_execution: Callable[[Plan, PlanRun, Step], BeforeStepExecutionOutcome] | None = None
+    before_step_execution: Callable[[PlanV2, PlanRunV2, StepV2], BeforeStepExecutionOutcome] | None = None
     """Called before executing each step.
 
     Args:
-        plan: The plan being executed
-        plan_run: The current plan run
-        step: The step about to be executed
+        plan: The PlanV2 being executed
+        plan_run: The current PlanRunV2
+        step: The StepV2 about to be executed
 
     Returns:
         BeforeStepExecutionOutcome | None: Whether to continue with the step execution or skip it.
             If None is returned, the default behaviour is to continue with the step execution.
     """
 
-    after_step_execution: Callable[[Plan, PlanRun, Step, Output], None] | None = None
+    after_step_execution: Callable[[PlanV2, PlanRunV2, StepV2, Output], None] | None = None
     """Called after executing each step.
 
     When there's an error, this is called with the error as the output value.
 
     Args:
-        plan: The plan being executed
-        plan_run: The current plan run
-        step: The step that was executed
+        plan: The PlanV2 being executed
+        plan_run: The current PlanRunV2
+        step: The StepV2 that was executed
         output: The output from the step execution
     """
 
-    before_plan_run: Callable[[Plan, PlanRun], None] | None = None
+    before_plan_run: Callable[[PlanV2, PlanRunV2], None] | None = None
     """Called before executing the first step of the plan run.
 
     Args:
-        plan: The plan being executed
-        plan_run: The current plan run
+        plan: The PlanV2 being executed
+        plan_run: The current PlanRunV2
     """
 
-    after_plan_run: Callable[[Plan, PlanRun, Output], None] | None = None
+    after_plan_run: Callable[[PlanV2, PlanRunV2, Output], None] | None = None
     """Called after executing the plan run.
 
     This is not called if a clarification is raised, as it is expected that the plan
     will be resumed after the clarification is handled.
 
     Args:
-        plan: The plan that was executed
-        plan_run: The completed plan run
+        plan: The PlanV2 that was executed
+        plan_run: The completed PlanRunV2
         output: The final output from the plan execution
     """
 
     before_tool_call: (
-        Callable[[Tool, dict[str, Any], PlanRun, Step], Clarification | None] | None
+        Callable[[Tool, dict[str, Any], PlanRunV2, StepV2], Clarification | None] | None
     ) = None
     """Called before the tool is called.
 
     Args:
         tool: The tool about to be called
         args: The args for the tool call. These are mutable and so can be modified in place as
           required.
-        plan_run: The current plan run
-        step: The step being executed
+        plan_run: The current PlanRunV2
+        step: The StepV2 being executed
 
     Returns:
         Clarification | None: A clarification to raise, or None to proceed with the tool call
     """
 
-    after_tool_call: Callable[[Tool, Any, PlanRun, Step], Clarification | None] | None = None
+    after_tool_call: Callable[[Tool, Any, PlanRunV2, StepV2], Clarification | None] | None = None
     """Called after the tool is called.
 
     Args:
         tool: The tool that was called
         output: The output returned from the tool call
-        plan_run: The current plan run
-        step: The step being executed
+        plan_run: The current PlanRunV2
+        step: The StepV2 being executed
 
     Returns:
         Clarification | None: A clarification to raise, or None to proceed. If a clarification
@@ -135,8 +136,8 @@
 def clarify_on_all_tool_calls(
     tool: Tool,
     args: dict[str, Any],
-    plan_run: PlanRun,
-    step: Step,
+    plan_run: PlanRunV2,
+    step: StepV2,
 ) -> Clarification | None:
     """Raise a clarification to check the user is happy with all tool calls before proceeding.
 
@@ -152,7 +153,7 @@
 
 def clarify_on_tool_calls(
     tool: str | Tool | list[str] | list[Tool],
-) -> Callable[[Tool, dict[str, Any], PlanRun, Step], Clarification | None]:
+) -> Callable[[Tool, dict[str, Any], PlanRunV2, StepV2], Clarification | None]:
     """Return a hook that raises a clarification before calls to the specified tool.
 
     Args:
@@ -185,8 +186,8 @@
 def _clarify_on_tool_call_hook(
     tool: Tool,
     args: dict[str, Any],
-    plan_run: PlanRun,
-    step: Step,  # noqa: ARG001
+    plan_run: PlanRunV2,
+    step: StepV2,  # noqa: ARG001
     tool_ids: list[str] | None,
 ) -> Clarification | None:
     """Raise a clarification to check the user is happy with all tool calls before proceeding."""
@@ -217,8 +218,8 @@
     return None
 
 
-def log_step_outputs(plan: Plan, plan_run: PlanRun, step: Step, output: Output) -> None:  # noqa: ARG001
+def log_step_outputs(plan: PlanV2, plan_run: PlanRunV2, step: StepV2, output: Output) -> None:  # noqa: ARG001
     """Log the output of a step in the plan."""
     logger().info(
-        f"Step with task {step.task} using tool {step.tool_id} completed with result: {output}"
+        f"Step {step.step_name} completed with result: {output}"
     )

portia/open_source_tools/browser_tool.py

@@ -615,10 +615,18 @@ def __init__(self, api_key: str | None = None, project_id: str | None = None) ->
         def step_complete(self, ctx: ToolRunContext) -> None:
             """Call when the step is complete closes the session to persist context."""
             # Only clean up the session on the final browser tool call
-            if any(
-                step.tool_id == "browser_tool"
-                for step in ctx.plan.steps[ctx.plan_run.current_step_index + 1 :]
-            ):
+            # Check if there are more browser tool steps coming up
+            from portia.builder.invoke_tool_step import InvokeToolStep
+
+            has_more_browser_steps = False
+            for step in ctx.plan.steps[ctx.plan_run.current_step_index + 1 :]:
+                if isinstance(step, InvokeToolStep):
+                    tool_id = step.tool if isinstance(step.tool, str) else step.tool.id
+                    if tool_id == "browser_tool":
+                        has_more_browser_steps = True
+                        break
+
+            if has_more_browser_steps:
                 return
 
             session_id = ctx.end_user.get_additional_data("bb_session_id")

portia/plan_run.py

@@ -17,6 +17,8 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, Any
+
 from pydantic import BaseModel, ConfigDict, Field
 
 from portia.clarification import (
@@ -25,9 +27,14 @@
     ClarificationListType,
 )
 from portia.common import PortiaEnum
+from portia.config import Config
+from portia.end_user import EndUser
 from portia.execution_agents.output import LocalDataValue, Output
 from portia.prefixed_uuid import PlanRunUUID, PlanUUID
 
+if TYPE_CHECKING:
+    from portia.builder.plan_v2 import PlanV2
+
 
 class PlanRunState(PortiaEnum):
     """The current state of the Plan Run.
@@ -228,3 +235,148 @@
             plan_run_inputs=plan_run.plan_run_inputs,
             structured_output_schema=plan_run.structured_output_schema,
         )
+
+
+class StepOutputValue(BaseModel):
+    """Value that can be referenced by name in a PlanV2 run.
+
+    Attributes:
+        value: The referenced value.
+        description: Description of the referenced value.
+        step_name: The name of the step that produced this value.
+        step_num: The step number that produced this value.
+
+    """
+
+    value: Any = Field(description="The referenced value.")
+    description: str = Field(description="Description of the referenced value.", default="")
+    step_name: str = Field(description="The name of the step that produced this value.")
+    step_num: int = Field(description="The step number that produced this value.")
+
+
+class PlanRunV2(BaseModel):
+    """A V2 plan run represents a running instance of a PlanV2.
+
+    This is the native representation for PlanV2 execution, replacing the legacy
+    PlanRun structure with a cleaner, more consistent data model.
+
+    Attributes:
+        id: A unique ID for this plan run.
+        state: The current state of the plan run.
+        current_step_index: The current step that is being executed.
+        plan: The PlanV2 being executed.
+        end_user: The end user executing the plan.
+        step_output_values: List of outputs from completed steps.
+        final_output: The final consolidated output of the plan run.
+        plan_run_inputs: Dict mapping plan input names to their values.
+        config: The Portia configuration.
+
+    """
+
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+
+    id: PlanRunUUID = Field(
+        default_factory=PlanRunUUID,
+        description="A unique ID for this plan run.",
+    )
+    state: PlanRunState = Field(
+        default=PlanRunState.NOT_STARTED,
+        description="The current state of the plan run.",
+    )
+    current_step_index: int = Field(
+        default=0,
+        description="The current step that is being executed.",
+    )
+    plan: PlanV2 = Field(
+        description="The PlanV2 being executed.",
+    )
+    end_user: EndUser = Field(
+        description="The end user executing the plan.",
+    )
+    step_output_values: list[StepOutputValue] = Field(
+        default_factory=list,
+        description="List of outputs from completed steps.",
+    )
+    final_output: Output | None = Field(
+        default=None,
+        description="The final consolidated output of the plan run.",
+    )
+    plan_run_inputs: dict[str, LocalDataValue] = Field(
+        default_factory=dict,
+        description="Dict mapping plan input names to their values.",
+    )
+    config: Config = Field(
+        description="The Portia configuration.",
+    )
+    clarifications: ClarificationListType = Field(
+        default_factory=list,
+        description="Clarifications raised during this plan run.",
+    )
+
+    def get_outstanding_clarifications(self) -> ClarificationListType:
+        """Return all outstanding clarifications.
+
+        Returns:
+            ClarificationListType: A list of outstanding clarifications that have not been resolved.
+
+        """
+        return [
+            clarification
+            for clarification in self.clarifications
+            if not clarification.resolved
+        ]
+
+    def get_clarifications_for_step(self, step: int | None = None) -> ClarificationListType:
+        """Return clarifications for the given step.
+
+        Args:
+            step: The step to get clarifications for. Defaults to current step.
+
+        Returns:
+            ClarificationListType: A list of clarifications for the given step.
+
+        """
+        if step is None:
+            step = self.current_step_index
+        return [
+            clarification
+            for clarification in self.clarifications
+            if clarification.step == step
+        ]
+
+    def get_clarification_for_step(
+        self, category: ClarificationCategory, step: int | None = None
+    ) -> Clarification | None:
+        """Return a clarification of the given category for the given step if it exists.
+
+        Args:
+            step: The step to get a clarification for. Defaults to current step.
+            category: The category of the clarification to get.
+
+        Returns:
+            Clarification | None: The clarification if found, None otherwise.
+
+        """
+        if step is None:
+            step = self.current_step_index
+        return next(
+            (
+                clarification
+                for clarification in self.clarifications
+                if clarification.step == step and clarification.category == category
+            ),
+            None,
+        )
+
+    def __str__(self) -> str:
+        """Return the string representation of the PlanRunV2.
+
+        Returns:
+            str: A string representation containing key run attributes.
+
+        """
+        return (
+            f"PlanRunV2(id={self.id}, plan_id={self.plan.id}, "
+            f"state={self.state}, current_step_index={self.current_step_index}, "
+            f"final_output={'set' if self.final_output else 'unset'})"
+        )