From 174ef162b3a79e9e684eaab87a0895fbba47cb36 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Wed, 29 Oct 2025 19:02:32 -0700
Subject: [PATCH] fix(mypy): add fast and full mypy modes (#3975)

`mypy` became very slow for the common path. This can make local
pre-commit runs very slow. Let's restore that.

- restore fast mirrors-mypy hook for local runs
- add optional mypy-full hook and docs so devs can match CI
- run full mypy in CI with a hint when failures occur

### Test Plan
- uv run pre-commit run mypy --all-files
- uv run pre-commit run mypy-full --hook-stage manual --all-files
- uv run --group dev --group type_checking mypy
---
 .github/workflows/pre-commit.yml              | 18 +++++++++++++-
 .pre-commit-config.yaml                       | 24 +++++++++++++------
 CONTRIBUTING.md                               | 12 ++++++++++
 .../agents/meta_reference/agent_instance.py   | 19 ++++++++-------
 4 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index 99ef87196..485009578 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -43,6 +43,9 @@ jobs:
           cache: 'npm'
           cache-dependency-path: 'src/llama_stack/ui/'
 
+      - name: Set up uv
+        uses: astral-sh/setup-uv@2ddd2b9cb38ad8efd50337e8ab201519a34c9f24 # v7.1.1
+
       - name: Install npm dependencies
         run: npm ci
         working-directory: src/llama_stack/ui
@@ -52,7 +55,7 @@ jobs:
         uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
         continue-on-error: true
         env:
-          SKIP: no-commit-to-branch
+          SKIP: no-commit-to-branch,mypy
           RUFF_OUTPUT_FORMAT: github
 
       - name: Check pre-commit results
@@ -109,3 +112,16 @@ jobs:
             echo "$unstaged_files"
             exit 1
           fi
+
+      - name: Sync dev + type_checking dependencies
+        run: uv sync --group dev --group type_checking
+
+      - name: Run mypy (full type_checking)
+        run: |
+          set +e
+          uv run --group dev --group type_checking mypy
+          status=$?
+          if [ $status -ne 0 ]; then
+            echo "::error::Full mypy failed. Reproduce locally with 'uv run pre-commit run mypy-full --hook-stage manual --all-files'."
+          fi
+          exit $status
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index f4d92403e..9990b6342 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -57,17 +57,27 @@ repos:
     hooks:
     -   id: uv-lock
 
--   repo: local
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.18.2
     hooks:
     -   id: mypy
-        name: mypy
         additional_dependencies:
-          - uv==0.7.8
-        entry: uv run --group dev --group type_checking mypy
-        language: python
-        types: [python]
+          - uv==0.6.2
+          - pytest
+          - rich
+          - types-requests
+          - pydantic
+          - httpx
         pass_filenames: false
-        require_serial: true
+
+-   repo: local
+    hooks:
+    -   id: mypy-full
+        name: mypy (full type_checking)
+        entry: uv run --group dev --group type_checking mypy
+        language: system
+        pass_filenames: false
+        stages: [manual]
 
 # - repo: https://github.com/tcort/markdown-link-check
 #   rev: v3.11.2
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index c869b4f5c..d84332829 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -61,6 +61,18 @@ uv run pre-commit run --all-files -v
 
 The `-v` (verbose) parameter is optional but often helpful for getting more information about any issues with that the pre-commit checks identify.
 
+To run the expanded mypy configuration that CI enforces, use:
+
+```bash
+uv run pre-commit run mypy-full --hook-stage manual --all-files
+```
+
+or invoke mypy directly with all optional dependencies:
+
+```bash
+uv run --group dev --group type_checking mypy
+```
+
 ```{caution}
 Before pushing your changes, make sure that the pre-commit hooks have passed successfully.
 ```
diff --git a/src/llama_stack/providers/inline/agents/meta_reference/agent_instance.py b/src/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
index b6fad553a..46619087c 100644
--- a/src/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
+++ b/src/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
@@ -995,14 +995,17 @@ class ChatAgent(ShieldRunnerMixin):
         except json.JSONDecodeError as e:
             raise ValueError(f"Failed to parse arguments for tool call: {tool_call.arguments}") from e
 
-        result = await self.tool_runtime_api.invoke_tool(
-            tool_name=tool_name_str,
-            kwargs={
-                "session_id": session_id,
-                # get the arguments generated by the model and augment with toolgroup arg overrides for the agent
-                **args,
-                **self.tool_name_to_args.get(tool_name_str, {}),
-            },
+        result = cast(
+            ToolInvocationResult,
+            await self.tool_runtime_api.invoke_tool(
+                tool_name=tool_name_str,
+                kwargs={
+                    "session_id": session_id,
+                    # get the arguments generated by the model and augment with toolgroup arg overrides for the agent
+                    **args,
+                    **self.tool_name_to_args.get(tool_name_str, {}),
+                },
+            ),
         )
         logger.debug(f"tool call {tool_name_str} completed with result: {result}")
         return result