Add AI integration testing suite

.github/workflows/ai-integration.yaml

@@ -0,0 +1,61 @@
+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; live-model tests require provider keys and are skipped on pull
+# requests from forks (which don't get secrets). Together they catch durable
+# replay and provider API/type breakage.
+on:
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  ai-integration:
+    name: "AI integration (scripted, Python 3.12)"
+    runs-on: warp-ubuntu-latest-x64-4x
+    env:
+      UV_PYTHON: "3.12"
+    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
+      # 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.12)"
+    runs-on: warp-ubuntu-latest-x64-4x
+    env:
+      UV_PYTHON: "3.12"
+      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

.github/workflows/ai-sdk-bump.yaml

@@ -0,0 +1,116 @@
+name: AI SDK Bump Test
+
+# Manually test a candidate agent SDK version before bumping it for real. The
+# direct SDK is pinned to the requested version, its dependency family is
+# re-resolved, and its Restate integration suite is run against that result.
+on:
+  workflow_dispatch:
+    inputs:
+      agent_sdk:
+        description: "Agent SDK to upgrade"
+        required: true
+        type: choice
+        options:
+          - openai-agents
+          - pydantic-ai
+          - langchain
+          - google-adk
+      version:
+        description: "Release version or tag (for example 0.7.0 or v0.7.0)"
+        required: true
+        type: string
+
+permissions:
+  contents: read
+
+jobs:
+  bump-and-test:
+    if: github.repository_owner == 'restatedev'
+    name: "Test ${{ inputs.agent_sdk }} ${{ inputs.version }}"
+    runs-on: warp-ubuntu-latest-x64-4x
+    env:
+      UV_PYTHON: "3.12"
+      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: Upgrade ${{ inputs.agent_sdk }} to ${{ inputs.version }}
+        id: upgrade
+        env:
+          AGENT_SDK: ${{ inputs.agent_sdk }}
+          VER: ${{ inputs.version }}
+        run: |
+          set -euo pipefail
+          VER="${VER#v}"
+          if [[ -z "$VER" ]]; then
+            echo "::error::version must not be empty"
+            exit 1
+          fi
+
+          case "$AGENT_SDK" in
+            openai-agents)
+              EXTRA=openai
+              SPEC=openai-agents
+              TEST_RECIPE=test-ai-openai
+              PACKAGES=(openai-agents)
+              ;;
+            pydantic-ai)
+              EXTRA=pydantic_ai
+              SPEC='pydantic-ai-slim[mcp,openai]'
+              TEST_RECIPE=test-ai-pydantic
+              PACKAGES=(pydantic-ai-slim)
+              ;;
+            google-adk)
+              EXTRA=adk
+              SPEC=google-adk
+              TEST_RECIPE=test-ai-google-adk
+              PACKAGES=(google-adk)
+              ;;
+            langchain)
+              EXTRA=langchain
+              SPEC=langchain
+              TEST_RECIPE=test-ai-langchain
+              PACKAGES=(langchain langchain-core langchain-openai langgraph)
+              ;;
+            *)
+              echo "::error::unknown agent SDK '$AGENT_SDK'"
+              exit 1
+              ;;
+          esac
+
+          echo "Pinning $SPEC==$VER in extra [$EXTRA]"
+          uv add --no-sync --optional "$EXTRA" "${SPEC}==${VER}"
+
+          upgrade_args=()
+          for package in "${PACKAGES[@]}"; do
+            upgrade_args+=(--upgrade-package "$package")
+          done
+          uv lock "${upgrade_args[@]}"
+
+          echo "test_recipe=$TEST_RECIPE" >> "$GITHUB_OUTPUT"
+          printf 'packages=%s\n' "${PACKAGES[*]}" >> "$GITHUB_OUTPUT"
+
+      - name: Install dependencies
+        run: just sync
+
+      - name: Show resolved version
+        env:
+          PACKAGES: ${{ steps.upgrade.outputs.packages }}
+        run: uv pip show $PACKAGES
+
+      - name: Build
+        run: just build
+
+      # KEY: fails on any changed types/signatures the integration code relies on.
+      - name: Type check
+        run: just typecheck
+
+      - name: Run ${{ inputs.agent_sdk }} integration tests (scripted + live)
+        env:
+          TEST_RECIPE: ${{ steps.upgrade.outputs.test_recipe }}
+        run: just "$TEST_RECIPE"

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`.