fixup! feat(tools): implement extended tool definitions with native PydanticAI integration

Author: Leoyzen

src/agentpool/agents/native_agent/tool_wrapping.py

@@ -150,7 +150,7 @@ async def wrapped(  # pyright: ignore[reportRedeclaration]
             if result == "allow":
                 # Populate AgentContext with RunContext data if needed
                 if agent_ctx.data is None:
-                    agent_ctx.data = ctx.deps
+                    agent_ctx.data = ctx.deps.data if ctx.deps else ctx.deps
 
                 if agent_ctx_key:  # inject AgentContext
                     # Build model_name from RunContext's model (provider:model_name format)

src/agentpool/tools/base.py

@@ -245,7 +245,18 @@ def _get_json_schema(self, func: Callable[..., Any] | None = None) -> dict[str,
             # Return only the parameters part (the "object" schema)
             # Use model_dump - schemez.FunctionSchema has this method (pydantic-compatible)
             schema_dump = getattr(schema, "model_dump")()  # noqa: B009, type: ignore[attr-defined]
-            return schema_dump["parameters"]  # type: ignore[no-any-return]
+            generated_params = schema_dump["parameters"]
+
+            # Apply parameter overrides to maintain consistency with the primary path
+            if "parameters" in self.schema_override:
+                override_params = self.schema_override["parameters"]
+                if "properties" in override_params:
+                    for param_name, param_def in override_params["properties"].items():
+                        if param_name in generated_params.get("properties", {}):
+                            generated_params["properties"][param_name].update(param_def)
+                        else:
+                            generated_params.setdefault("properties", {})[param_name] = param_def
+            return generated_params  # type: ignore[no-any-return]
         else:
             return schema.json_schema
 

tests/tools/test_runcontext.py

@@ -27,11 +27,7 @@ async def data_with_run_ctx(ctx: RunContext[AgentContext[dict[str, str]]]) -> st
 
 async def data_with_agent_ctx(ctx: AgentContext) -> str:
     """Tool accessing data through AgentContext."""
-    # When a tool requests AgentContext, it gets RunContext.deps
-    # RunContext.deps is AgentContext, and the user data is in AgentContext.data
-    # But ctx here is AgentContext, so ctx.data contains the user data directly
-    data_value = ctx.data.data if isinstance(ctx.data, AgentContext) else ctx.data
-    return f"Data from AgentContext: {data_value}"
+    return f"Data from AgentContext: {ctx.data}"
 
 
 async def no_ctx_tool() -> str: