Bugfix: avoid race condition when refreshing google token

[JadHADDAD92]

Hello,
I have been having an issue with the refresh token when launching multiple Gemini agents concurrently.
Here is a minimal reproducible code:

Class handling agent

class InformationExtractionAgent:
    def __init__(self):
        self.model = GeminiModel(
            "gemini-2.5-pro",
            provider=GoogleVertexProvider(
                service_account_file=os.path.join("shared", "config", "gemini.json"),
                region="europe-west1",
            ),
        )
        self.settings = GeminiModelSettings(
            gemini_thinking_config={"thinking_budget": 128, "include_thoughts": True}
        )

        self.agent = Agent(
            model=self.model,
            model_settings=self.settings,
            output_type=DocumentExtractionResult,
            system_prompt="..."
        )

    async def extract_information(...) -> AgentRunResult[DocumentExtractionResult]:
        result = await self.agent.run("...")
        return result

This class is called and used from a coroutine

async def extraction_process(...):
    information_extraction_agent = InformationExtractionAgent()
    # [...]
    extraction_result = await information_extraction_agent.extract_information(...)
    # [...]

and this coroutine is called concurrently in a Semaphore

async run_extraction_job(tasks_ids):
    semaphore = Semaphore(10)
    # [...]
    async def semaphore_execute_task(task_id: int):
        async with semaphore:
             await extraction_process(task_id)
    # [...]
    await asyncio.gather(*[semaphore_execute_task(task_id) for task_id in task_ids])

I would randomly get this error, not consistently, even though all the instances have the service_account_file configured to shared/config/gemini.json, seems like a race condition:

backend-1  |   File "/usr/local/lib/python3.11/site-packages/httpx/_client.py", line 1654, in _send_handling_auth
backend-1  |     request = await auth_flow.__anext__()
backend-1  |               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
backend-1  |   File "/usr/local/lib/python3.11/site-packages/pydantic_ai/providers/google_vertex.py", line 134, in async_auth_flow
backend-1  |     await self._refresh_token()
backend-1  |   File "/usr/local/lib/python3.11/site-packages/pydantic_ai/providers/google_vertex.py", line 169, in _refresh_token
backend-1  |     assert isinstance(self.credentials.token, str), f'Expected token to be a string, got {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
backend-1  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
backend-1  | AssertionError: Expected token to be a string, got None

with the help of Mert Sırakaya, here is a patch for this error, I have tried running the test many times after the patch, I don't have any errors, and then tried again without the patch and I occasionally got the error again.

For the typing lint I have put ignore[reportUnknownMemberType] comments because I didn't know what else to do, I'm open to better suggestions

[DouweM]

@JadHADDAD92 Thanks Jad. This look like the same issue as #1440. There, my colleague Marcelo this should be solved with the new GoogleModel and GoogleProvider(vertexai=True) you can use instead of GeminiModel and GoogleVertexProvider as in the example here: https://ai.pydantic.dev/models/google/#vertex-ai-enterprisecloud. Can you see if that works?

Note that this PR may still make sense, but I want to verify the new classes also work for you.

[JadHADDAD92]

I have tried GoogleModel with the following configs

class InformationExtractionAgent:
    def __init__(self):
        credentials = service_account.Credentials.from_service_account_file(
            os.path.join("shared", "config", "gemini.json"),
            scopes=["https://www.googleapis.com/auth/cloud-platform"],
        )
        self.model = GoogleModel(
            "gemini-2.5-pro",
            provider=GoogleProvider(
                credentials=credentials,
                location="europe-west1",
            ),
        )
        self.settings = GoogleModelSettings(
            google_thinking_config={"thinking_budget": 128, "include_thoughts": True}
        )

And couldn't reproduce my error, so I guess it works!

[DouweM]

@JadHADDAD92 Glad to hear it!

@Kludex What do you think, should we keep fixing issues in GeminiModel or deprecate it in v1 and push people to GoogleModel instead?