Add AI integration testing suite

.github/workflows/ai-integration.yaml

@@ -0,0 +1,82 @@
+name: AI Integration
+
+# Scripted and real-model integration tests for the AI SDK integrations
+# (openai-agents, pydantic-ai, Google ADK, and LangChain). Scripted tests run on
+# every PR against the newest allowed SDK versions outside the dependency
+# cooldown. Live-model tests use the committed lockfile, require provider keys,
+# and are skipped on pull requests from forks. Together they catch durable
+# replay and provider API/type breakage without exposing secrets to new packages.
+on:
+  pull_request:
+    branches: [main]
+  schedule:
+    - cron: "17 5 * * 1"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+env:
+  # Reject dependencies reported as malicious by OSV/OpenSSF during sync.
+  UV_MALWARE_CHECK: "1"
+
+jobs:
+  ai-integration:
+    name: "AI integration (scripted, Python 3.14)"
+    runs-on: warp-ubuntu-latest-x64-4x
+    env:
+      UV_PYTHON: "3.14"
+    steps:
+      - uses: actions/checkout@v4
+      - uses: extractions/setup-just@v2
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+      - name: Resolve the newest allowed AI SDKs
+        run: >-
+          uv lock
+          --upgrade-package google-adk
+          --upgrade-package google-genai
+          --upgrade-package openai
+          --upgrade-package openai-agents
+          --upgrade-package pydantic-ai-slim
+          --upgrade-package pydantic-graph
+          --upgrade-package langchain
+          --upgrade-package langchain-core
+          --upgrade-package langchain-openai
+          --upgrade-package langgraph
+          --upgrade-package langgraph-prebuilt
+      - name: Install dependencies
+        run: just sync
+      - name: Build
+        run: just build
+      # Scripted only here: this job has no provider keys and runs on forks.
+      - name: Run scripted AI integration tests
+        run: just test-ai-scripted
+
+  ai-integration-live:
+    if: >-
+      github.repository_owner == 'restatedev' &&
+      (github.event_name == 'workflow_dispatch' ||
+       (github.event_name == 'pull_request' &&
+        github.event.pull_request.head.repo.full_name == github.repository &&
+        github.event.pull_request.user.login != 'dependabot[bot]'))
+    name: "AI integration (live model, Python 3.14)"
+    runs-on: warp-ubuntu-latest-x64-4x
+    env:
+      UV_PYTHON: "3.14"
+      OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+      GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: extractions/setup-just@v2
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+      - name: Install dependencies
+        run: just sync
+      - name: Build
+        run: just build
+      # Full suite (scripted + live) -- this job has the key.
+      - name: Run AI integration tests
+        run: just test-ai

Justfile

@@ -21,17 +21,40 @@ typecheck-pyright:
     PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright examples/ 
     PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright tests
     PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright test-services/
+    PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright ai-tests/
 
 typecheck-mypy:
     uv run -m mypy --check-untyped-defs --ignore-missing-imports --implicit-optional python/
     uv run -m mypy --check-untyped-defs --ignore-missing-imports --implicit-optional examples/
     uv run -m mypy --check-untyped-defs --ignore-missing-imports --implicit-optional tests/
+    uv run -m mypy --check-untyped-defs --ignore-missing-imports --implicit-optional ai-tests/
 
 typecheck: typecheck-pyright typecheck-mypy
 
 test:
     uv run -m pytest tests/*
 
+# Each integration runs in its own dependency environment.
+test-ai-openai:
+    uv run --isolated --locked --extra test --extra harness --extra serde --extra openai -m pytest ai-tests/openai_test.py -v
+
+test-ai-pydantic:
+    uv run --isolated --locked --extra test --extra harness --extra serde --extra pydantic_ai -m pytest ai-tests/pydantic_test.py -v
+
+test-ai-google-adk:
+    uv run --isolated --locked --extra test --extra harness --extra serde --extra adk -m pytest ai-tests/google_adk_test.py -v
+
+test-ai-langchain:
+    uv run --isolated --locked --extra test --extra harness --extra serde --extra langchain --extra langchain_test -m pytest ai-tests/langchain_test.py -v
+
+test-ai: test-ai-openai test-ai-pydantic test-ai-google-adk test-ai-langchain
+
+test-ai-scripted:
+    uv run --isolated --locked --extra test --extra harness --extra serde --extra openai -m pytest ai-tests/openai_test.py -m "not live_model" -v
+    uv run --isolated --locked --extra test --extra harness --extra serde --extra pydantic_ai -m pytest ai-tests/pydantic_test.py -m "not live_model" -v
+    uv run --isolated --locked --extra test --extra harness --extra serde --extra adk -m pytest ai-tests/google_adk_test.py -m "not live_model" -v
+    uv run --isolated --locked --extra test --extra harness --extra serde --extra langchain --extra langchain_test -m pytest ai-tests/langchain_test.py -m "not live_model" -v
+
 
 # Recipe to run both mypy and pylint
 verify: format lint typecheck test

ai-tests/README.md

@@ -0,0 +1,95 @@
+# AI integration tests
+
+End-to-end tests for the AI SDK integrations in `restate.ext.*`:
+
+- OpenAI Agents
+- Pydantic AI
+- Google ADK
+- LangChain
+
+The tests run each integration against a real Restate server. Their purpose is
+to catch breaking changes when either Restate or an upstream agent SDK is
+upgraded.
+
+## What we test
+
+Every scenario runs with `always_replay=True`, forcing Restate to suspend and
+replay on each await. This exercises the path most likely to expose:
+
+- journal mismatches and other non-determinism,
+- infinite failure or retry loops,
+- incompatible agent SDK response types,
+- interference between concurrent invocations or sessions.
+
+`disable_retries=True` makes ordinary handler failures surface immediately
+instead of hiding them behind retry backoff.
+
+The common scenarios cover:
+
+| Scenario | What it checks |
+| --- | --- |
+| Tool call | The model requests a tool and receives its output |
+| Multi-turn session | Conversation state survives across virtual object calls |
+| Concurrent invocations | Distinct object keys do not interfere with each other |
+| Parallel tools | Multiple tool calls pass through the integration turnstile |
+| Terminal tool error | A permanent failure surfaces without retrying forever |
+| Local handoff | Delegation between agents is serialized correctly |
+| Remote handoff | Agent state and tool results survive a durable RPC |
+
+These are integration tests, not model evaluations. They check that the agent
+SDK protocol remains compatible and deterministic, not whether an answer is
+subjectively good.
+
+## Model modes
+
+The suite uses two complementary modes:
+
+- **Scripted:** deterministic provider responses exercise the real agent SDK
+  types and integration code without credentials.
+- **Live:** small real OpenAI or Gemini calls catch upstream response type and
+  format changes that scripted responses may miss.
+
+All scenarios run in scripted mode. The tool-call and multi-turn scenarios also
+run in live mode.
+
+## Running the tests
+
+Docker is required.
+
+Run all deterministic tests without provider credentials:
+
+```shell
+just test-ai-scripted
+```
+
+Run all scripted and live tests:
+
+```shell
+export OPENAI_API_KEY=sk-...
+export GOOGLE_API_KEY=...
+just test-ai
+```
+
+Run one integration, including its live tests:
+
+```shell
+just test-ai-openai
+just test-ai-pydantic
+just test-ai-google-adk
+just test-ai-langchain
+```
+
+Each integration runs in an isolated, locked dependency environment. Live tests
+fail when their required key is missing; they do not silently skip.
+
+## CI
+
+The `AI Integration` workflow runs scripted tests on every pull request. Live
+tests run only for trusted pull requests where repository secrets are
+available.
+
+The manually triggered `AI SDK Bump Test` workflow accepts an agent SDK and a
+version, resolves that candidate dependency family, type-checks the integration,
+and runs its scripted and live tests.
+
+These tests are intentionally separate from `just test` and `just verify`.

ai-tests/google_adk_model_stub.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+
+from google.adk.models.google_llm import Gemini
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+
+CITIES = ("Paris", "London", "Tokyo", "Berlin", "Rome")
+
+
+async def scripted_generate_content_async(
+    model: Gemini,
+    llm_request: LlmRequest,
+    stream: bool = False,
+) -> AsyncGenerator[LlmResponse, None]:
+    """Replace only the Gemini provider call while retaining normal ADK setup."""
+    del model, stream
+    tool_names = set(llm_request.tools_dict)
+    prompt = _latest_prompt(llm_request)
+
+    if "In which country is it?" in prompt:
+        yield _text_response("France")
+        return
+
+    if "What is the capital of France?" in prompt:
+        yield _text_response("Paris")
+        return
+
+    if "get_weather" in tool_names:
+        cities = [city for city in CITIES if city in prompt]
+        if not cities:
+            yield _text_response("Done.")
+            return
+        call_ids = {f"weather-{city.lower()}" for city in cities}
+        if _has_function_responses(llm_request, call_ids):
+            yield _text_response("Done.")
+            return
+        yield LlmResponse(
+            content=types.Content(
+                role="model",
+                parts=[
+                    types.Part(
+                        function_call=types.FunctionCall(
+                            id=f"weather-{city.lower()}",
+                            name="get_weather",
+                            args={"city": city},
+                        )
+                    )
+                    for city in cities
+                ],
+            )
+        )
+        return
+
+    if "explode" in tool_names:
+        yield _tool_call("explode", {"reason": "scripted failure"}, "explode")
+        return
+
+    if "ask_specialist" in tool_names:
+        if _has_function_responses(llm_request, {"ask-specialist"}):
+            yield _text_response("Done.")
+        else:
+            yield _tool_call(
+                "ask_specialist",
+                {"question": "How do I speed up a slow SQL query?"},
+                "ask-specialist",
+            )
+        return
+
+    if "transfer_to_agent" in tool_names:
+        if _has_function_responses(llm_request, {"billing-handoff"}):
+            yield _text_response("Done.")
+        else:
+            yield _tool_call(
+                "transfer_to_agent",
+                {"agent_name": "billing_agent"},
+                "billing-handoff",
+            )
+        return
+
+    yield _text_response("Done.")
+
+
+def _latest_prompt(llm_request: LlmRequest) -> str:
+    for content in reversed(llm_request.contents):
+        if content.role != "user":
+            continue
+        for part in reversed(content.parts or []):
+            if part.text:
+                return part.text
+    return ""
+
+
+def _has_function_responses(llm_request: LlmRequest, call_ids: set[str]) -> bool:
+    return any(
+        part.function_response is not None and part.function_response.id in call_ids
+        for content in llm_request.contents
+        for part in content.parts or []
+    )
+
+
+def _text_response(text: str) -> LlmResponse:
+    return LlmResponse(content=types.Content(role="model", parts=[types.Part.from_text(text=text)]))
+
+
+def _tool_call(name: str, args: dict[str, str], call_id: str) -> LlmResponse:
+    return LlmResponse(
+        content=types.Content(
+            role="model",
+            parts=[types.Part(function_call=types.FunctionCall(id=call_id, name=name, args=args))],
+        )
+    )