diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000000..3326877816 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,88 @@ +version: 2 + +updates: + # Python. `uv` rather than `pip`: it resolves through the uv CLI and rewrites + # pyproject.toml alongside uv.lock, so the declared pins here — which carry + # policy, not just versions — do not drift away from the lock file. + - package-ecosystem: uv + directory: / + schedule: + interval: weekly + day: monday + open-pull-requests-limit: 5 + groups: + python-minor-and-patch: + applies-to: version-updates + patterns: + - "*" + update-types: + - minor + - patch + python-major: + applies-to: version-updates + patterns: + - "*" + update-types: + - major + commit-message: + prefix: chore(deps) + + # GitHub Actions. Every `uses:` here is `@ # vX.Y.Z` on one line, which + # is the format Dependabot rewrites, comment included. + - package-ecosystem: github-actions + directory: / + schedule: + interval: weekly + day: monday + open-pull-requests-limit: 5 + groups: + actions-minor-and-patch: + applies-to: version-updates + patterns: + - "*" + update-types: + - minor + - patch + actions-major: + applies-to: version-updates + patterns: + - "*" + update-types: + - major + ignore: + # .github/workflows/pr-agent-review.yaml pins this action to a commit on + # `main`, not to a release. There is no version to bump to, and the + # documented behaviour for an untagged sha is to move it to whatever HEAD + # happens to be. Moving the self-review action is a deliberate act. + - dependency-name: Codium-ai/pr-agent + commit-message: + prefix: chore(ci) + + # Base images. The docker file fetcher matches any file whose name contains + # "dockerfile", so `/` picks up Dockerfile.github_action_dockerhub and + # `/docker` picks up both Dockerfile and Dockerfile.lambda. + - package-ecosystem: docker + directories: + - / + - /docker + schedule: + interval: weekly + day: monday + open-pull-requests-limit: 3 + ignore: + # Dockerfile.github_action_dockerhub builds FROM this project's own + # published image, tracked by a stage tag rather than a version. + - dependency-name: pragent/pr-agent + commit-message: + prefix: chore(docker) + + # Hook revisions in .pre-commit-config.yaml. The `local` repo has no upstream + # and is left alone by the ecosystem. + - package-ecosystem: pre-commit + directory: / + schedule: + interval: weekly + day: monday + open-pull-requests-limit: 3 + commit-message: + prefix: chore(pre-commit) diff --git a/.github/workflows/build-and-test.yaml b/.github/workflows/build-and-test.yaml index 16f76f89ed..4f54e6898e 100644 --- a/.github/workflows/build-and-test.yaml +++ b/.github/workflows/build-and-test.yaml @@ -14,15 +14,21 @@ jobs: steps: - id: checkout - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - id: uv + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + + - name: Check uv.lock + run: uv lock --check - id: dockerx name: Setup Docker Buildx - uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0 + uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0 - id: build name: Build dev docker - uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6.19.2 + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0 with: context: . file: ./docker/Dockerfile @@ -47,7 +53,7 @@ jobs: exit 1 fi - name: Upload coverage to Codecov - uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4 + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 with: token: ${{ secrets.CODECOV_TOKEN }} @@ -55,7 +61,7 @@ jobs: # Smoke-start github_app (gunicorn + uvicorn) with no secrets; GET / → 200. - id: build-smoke name: Build github_app docker - uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6.19.2 + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0 with: context: . file: ./docker/Dockerfile diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 303d59d7da..dec3c7bf80 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -31,14 +31,14 @@ jobs: - language: python build-mode: none steps: - - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - - uses: github/codeql-action/init@7fd177fa680c9881b53cdab4d346d32574c9f7f4 # v3.35.4 + - uses: github/codeql-action/init@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 with: languages: ${{ matrix.language }} build-mode: ${{ matrix.build-mode }} queries: security-extended,security-and-quality - - uses: github/codeql-action/analyze@7fd177fa680c9881b53cdab4d346d32574c9f7f4 # v3.35.4 + - uses: github/codeql-action/analyze@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 with: category: "/language:${{ matrix.language }}" diff --git a/.github/workflows/docs-ci.yaml b/.github/workflows/docs-ci.yaml index fb1056c4eb..946daac080 100644 --- a/.github/workflows/docs-ci.yaml +++ b/.github/workflows/docs-ci.yaml @@ -12,16 +12,16 @@ jobs: deploy: runs-on: ubuntu-latest steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Configure Git Credentials run: | git config user.name github-actions[bot] git config user.email 41898282+github-actions[bot]@users.noreply.github.com - - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: 3.x - run: echo "cache_id=$(date --utc '+%V')" >> "$GITHUB_ENV" - - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: key: mkdocs-material-${{ env.cache_id }} path: .cache diff --git a/.github/workflows/e2e_tests.yaml b/.github/workflows/e2e_tests.yaml index 0d33afb903..77255e4b7e 100644 --- a/.github/workflows/e2e_tests.yaml +++ b/.github/workflows/e2e_tests.yaml @@ -11,14 +11,14 @@ jobs: name: PR-Agent E2E GitHub App Test steps: - name: Checkout repository - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Docker Buildx - uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0 + uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0 - id: build name: Build dev docker - uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6.19.2 + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0 with: context: . file: ./docker/Dockerfile diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 24aa2de365..0cafc1f026 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -1,9 +1,10 @@ -# disabled. We might run it manually if needed. +# Runs every pre-commit hook (ruff, actionlint, YAML/TOML checks, ...) on +# pull requests, so the hooks are enforced in CI rather than manual-only. name: pre-commit on: workflow_dispatch: -# pull_request: + pull_request: # push: # branches: [main] @@ -11,7 +12,8 @@ jobs: pre-commit: runs-on: ubuntu-latest steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 # SEE https://github.com/pre-commit/action - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 diff --git a/.github/workflows/provider-smoke.yaml b/.github/workflows/provider-smoke.yaml new file mode 100644 index 0000000000..a4a5546132 --- /dev/null +++ b/.github/workflows/provider-smoke.yaml @@ -0,0 +1,32 @@ +name: Provider smoke + +on: + pull_request: + branches: + - main + push: + branches: + - main + +permissions: + contents: read + pull-requests: read + issues: read + +jobs: + github-provider-smoke: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Setup uv + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + + - name: Install dependencies + run: uv sync + + - name: Smoke /describe on a real GitHub PR + env: + GITHUB.USER_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: PYTHONPATH=. uv run pytest tests/health_test/test_github_smoke.py -q diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 74990eb7a2..7e6bd0d7b9 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -83,11 +83,11 @@ jobs: runs-on: ubuntu-latest environment: release steps: - - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ needs.prepare.outputs.sha }} - - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: '3.12' @@ -99,7 +99,7 @@ jobs: python -m pip install --upgrade pip build python -m build - - uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 + - uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2 with: password: ${{ secrets.PYPI_API_TOKEN }} # Re-running the workflow for the same version (e.g. after a @@ -168,25 +168,25 @@ jobs: suffix: gitlab_lambda rolling: gitlab_lambda steps: - - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ needs.prepare.outputs.sha }} - - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: '3.12' - name: Set version in pyproject.toml run: python scripts/set_pyproject_version.py "${{ needs.prepare.outputs.version }}" - - uses: docker/setup-qemu-action@c7c53464625b32c7a7e944ae62b3e17d2b600130 # v3.7.0 + - uses: docker/setup-qemu-action@1f40c72289eff860ee54a304f1438e3cff362e0a # v4.3.0 with: # amd64 is the runner's native arch; only arm64 needs emulation. platforms: arm64 - - uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0 + - uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0 - - uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.7.0 + - uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4.6.0 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} @@ -208,7 +208,7 @@ jobs: echo "tags=$TAGS" >> "$GITHUB_OUTPUT" - id: push - uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6.19.2 + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0 with: context: . file: ${{ matrix.dockerfile }} @@ -219,7 +219,7 @@ jobs: cache-from: type=gha,scope=${{ matrix.suffix || 'cli' }} cache-to: type=gha,mode=max,scope=${{ matrix.suffix || 'cli' }} - - uses: actions/attest-build-provenance@977bb373ede98d70efdf65b84cb5f73e068dcc2a # v3 + - uses: actions/attest-build-provenance@4d101475d8b20a2381f78447822ac1eab6504dd8 # v4.2.2 with: subject-name: index.docker.io/pragent/pr-agent subject-digest: ${{ steps.push.outputs.digest }} @@ -245,16 +245,18 @@ jobs: # the push rather than rebasing onto a newer commit that wasn't built. # This preserves the invariant that `git checkout vX.Y.Z` resolves to a # commit whose tree was published. - - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ needs.prepare.outputs.sha }} fetch-depth: 0 token: ${{ secrets.RELEASE_TOKEN || secrets.GITHUB_TOKEN }} - - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: '3.12' + - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + - name: Flag missing PyPI publish if: needs.publish-pypi.result != 'success' env: @@ -267,8 +269,10 @@ jobs: git config user.name "github-actions[bot]" git config user.email "41898282+github-actions[bot]@users.noreply.github.com" - - name: Bump pyproject.toml - run: python scripts/set_pyproject_version.py "${{ needs.prepare.outputs.version }}" + - name: Bump pyproject.toml and sync uv.lock + run: | + python scripts/set_pyproject_version.py "${{ needs.prepare.outputs.version }}" + uv lock - name: Commit & push bump (fast-forward only) id: commit @@ -276,19 +280,19 @@ jobs: VERSION: ${{ needs.prepare.outputs.version }} PUBLISHED_SHA: ${{ needs.prepare.outputs.sha }} run: | - if git diff --quiet pyproject.toml; then - echo "pyproject.toml already at $VERSION on $PUBLISHED_SHA; nothing to commit." + if git diff --quiet pyproject.toml uv.lock; then + echo "pyproject.toml and uv.lock already at $VERSION on $PUBLISHED_SHA; nothing to commit." echo "bumped=false" >> "$GITHUB_OUTPUT" echo "sha=$PUBLISHED_SHA" >> "$GITHUB_OUTPUT" exit 0 fi - git add pyproject.toml + git add pyproject.toml uv.lock git commit -m "chore(release): bump version to $VERSION [skip ci]" if git push origin "HEAD:refs/heads/main"; then echo "bumped=true" >> "$GITHUB_OUTPUT" echo "sha=$(git rev-parse HEAD)" >> "$GITHUB_OUTPUT" else - echo "::warning::main has moved past the published commit ($PUBLISHED_SHA); skipping the bump push rather than rebasing onto a commit that was not built. Bump pyproject.toml on main manually: git checkout main && python scripts/set_pyproject_version.py $VERSION && git commit -am 'chore(release): bump version to $VERSION' && git push" + echo "::warning::main has moved past the published commit ($PUBLISHED_SHA); skipping the bump push rather than rebasing onto a commit that was not built. Bump pyproject.toml and uv.lock on main manually: git checkout main && python scripts/set_pyproject_version.py $VERSION && uv lock && git commit -am 'chore(release): bump version to $VERSION' && git push" echo "bumped=false" >> "$GITHUB_OUTPUT" echo "sha=$PUBLISHED_SHA" >> "$GITHUB_OUTPUT" fi diff --git a/.github/workflows/release-drafter.yml b/.github/workflows/release-drafter.yml index b3cc0b2fcf..58be517edd 100644 --- a/.github/workflows/release-drafter.yml +++ b/.github/workflows/release-drafter.yml @@ -20,6 +20,6 @@ jobs: pull-requests: write runs-on: ubuntu-latest steps: - - uses: release-drafter/release-drafter@67e173cadb2fbd3de94f4a861e0c48c913b462ae # v6.4.0 + - uses: release-drafter/release-drafter@6a93d829887aa2e0748befe2e808c66c0ec6e4c7 # v6.4.0 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.gitignore b/.gitignore index 040f445e1a..3169689682 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,7 @@ dist/ build/ .DS_Store docs/.cache/ +docs/site/ .qodo poetry.lock .log diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7d6f90b468..47f9c6ff34 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -6,7 +6,7 @@ default_language_version: repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v6.0.0 + rev: 3e8a8703264a2f4a69428a0aa4dcb512790b2c8c # frozen: v6.0.0 hooks: - id: check-added-large-files exclude: ^uv\.lock$ # lock file, 686 KB and generated @@ -28,7 +28,7 @@ repos: exclude: ^tests/unittest/test_extend_patch\.py$ # trailing space; stripping it breaks the test. - repo: https://github.com/rhysd/actionlint - rev: v1.7.12 + rev: 914e7df21a07ef503a81201c76d2b11c789d3fca # frozen: v1.7.12 hooks: - id: actionlint @@ -41,11 +41,3 @@ repos: language: system types_or: [python, pyi] args: [--fix, --show-fixes] - - # - repo: https://github.com/PyCQA/bandit - # rev: 1.7.10 - # hooks: - # - id: bandit - # args: [ - # "-c", "pyproject.toml", - # ] diff --git a/AGENTS.md b/AGENTS.md index 02fab8b2b8..6139521989 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -76,7 +76,7 @@ Ruff is the single linting tool: `pyproject.toml` configures it and the pre-comm - Keep Python lines within the 120-character limit declared in `pyproject.toml`. - `pyproject.toml` configures Ruff rules `E`, `F`, `B` and `I`; `I001` (import sorting), `F401` (unused imports), and `F541` (f-strings without placeholders) are fixable. The `lint.ignore` list defers pre-existing violations — treat it as a debt ledger: fix the code and drop entries rather than adding new ones. - Before committing, run `uv run ruff check --fix` on the files you touched and fix every issue it reports. Keep fixes mechanical (rename, remove unused imports, sort imports); do not alter program logic while cleaning up — if a lint fix would change behavior, surface it instead of applying it silently. -- `.pre-commit-config.yaml` is the source of truth for enabled pre-commit hooks. Run them on the files you touched with `uv run pre-commit run --files ` and review the automatic edits so unrelated changes are not included. The pre-commit GitHub Actions workflow is manual-only (`workflow_dispatch`); the hooks are not enforced in CI. +- `.pre-commit-config.yaml` is the source of truth for enabled pre-commit hooks. Run them on the files you touched with `uv run pre-commit run --files ` and review the automatic edits so unrelated changes are not included. The pre-commit GitHub Actions workflow runs on pull requests, so the hooks are enforced in CI. - No general-purpose Python formatter is currently enforced (`ruff format` is deliberately not adopted yet). Preserve the surrounding file's formatting and avoid unrelated rewrites or repository-wide formatting. - Prefer double quotes for Python strings where consistent with the surrounding file. - Match existing docstring and comment style—concise English comments using imperative phrasing only where necessary. @@ -97,7 +97,7 @@ Ruff is the single linting tool: `pyproject.toml` configures it and the pre-comm - Follow `CONTRIBUTING.md`: keep changes focused, add or update tests, and use Conventional Commit-style messages (e.g., `fix: handle missing repo settings gracefully`). - Target branch names follow `feature/` or `fix/` patterns for substantial work. - Reference related issues and update README or docs when user-facing behavior shifts. -- Ensure CI workflows (`build-and-test`, `code-coverage`, `docs-ci`) succeed locally or in draft PRs before requesting review; reproduce failures with the documented commands above. +- Before requesting review, run the relevant local checks above and make sure the corresponding CI workflows (`build-and-test`, `pre-commit`) pass; coverage is collected and uploaded by `build-and-test`, and `docs-ci` runs only on pushes to `main` and `add-docs-portal`. - Include screenshots or terminal captures when modifying user-visible output or documentation previews. ## Safety and Permissions diff --git a/README.md b/README.md index 93b81f73df..0373ef7f60 100644 --- a/README.md +++ b/README.md @@ -106,25 +106,6 @@ pr-agent --pr_url https://github.com/owner/repo/pull/123 review Full notes for every release are on the [Releases page](https://github.com/the-pr-agent/pr-agent/releases). -### Jul 26, 2026 — [v0.41.0](https://github.com/the-pr-agent/pr-agent/releases/tag/v0.41.0) - -Claude Opus 5 support, and `docker/mosaico` became a self-contained deployment bundle. - -### Jul 25, 2026 — [v0.40.0](https://github.com/the-pr-agent/pr-agent/releases/tag/v0.40.0) - -Default model moved to GPT-5.6, Gemini 3.6 support, persistent inline comments (no more duplicate -suggestions across runs), an OpenRouter provider-routing/reasoning config, a tokenless -[plain-diff provider](https://docs.pr-agent.ai/usage-guide/plain_diff_mode/), and CI artifact -context injection. - -### Jul 5, 2026 — [v0.39.0](https://github.com/the-pr-agent/pr-agent/releases/tag/v0.39.0) - -`AGENTS.md` and friends are now fed to `/review`, `/describe` and `/improve` **by default**, so the -model picks up your project's conventions out of the box. Also: -[Agent Skills (`SKILL.md`)](https://docs.pr-agent.ai/core-abilities/agent_skills/), -[organization-level settings](https://docs.pr-agent.ai/usage-guide/configuration_options/), -[restricted mode](https://docs.pr-agent.ai/usage-guide/additional_configurations/#restricted-mode) -for reduced GitHub permissions, GitHub Checks as an output target, and Claude Sonnet 5 support. ## Why Use PR-Agent? diff --git a/docker/Dockerfile b/docker/Dockerfile index 11467cf521..a99e72561a 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,64 +1,64 @@ -FROM python:3.12.14-slim AS base +FROM python:3.14.7-slim AS base COPY --from=ghcr.io/astral-sh/uv:0.12.10 /uv /uvx /bin/ -RUN apt update && apt install --no-install-recommends -y git curl && apt-get clean && rm -rf /var/lib/apt/lists/* +RUN apt-get update && apt-get install --no-install-recommends -y git curl && apt-get clean && rm -rf /var/lib/apt/lists/* WORKDIR /app ENV UV_COMPILE_BYTECODE=1 -ADD pyproject.toml uv.lock ./ -ADD docs docs +COPY pyproject.toml uv.lock ./ +COPY docs docs RUN uv sync --frozen --no-install-project --no-dev ENV PATH="/app/.venv/bin:$PATH" ENV PYTHONPATH=/app FROM base AS github_app -ADD pr_agent pr_agent +COPY pr_agent pr_agent CMD ["python", "-m", "gunicorn", "-k", "uvicorn.workers.UvicornWorker", "-c", "pr_agent/servers/gunicorn_config.py", "--forwarded-allow-ips", "*", "pr_agent.servers.github_app:app"] FROM base AS bitbucket_app -ADD pr_agent pr_agent +COPY pr_agent pr_agent CMD ["python", "pr_agent/servers/bitbucket_app.py"] FROM base AS bitbucket_server_webhook -ADD pr_agent pr_agent +COPY pr_agent pr_agent CMD ["python", "pr_agent/servers/bitbucket_server_webhook.py"] FROM base AS github_polling -ADD pr_agent pr_agent +COPY pr_agent pr_agent CMD ["python", "pr_agent/servers/github_polling.py"] FROM base AS github_action -ADD pr_agent pr_agent -ADD github_action/entrypoint.sh / +COPY pr_agent pr_agent +COPY github_action/entrypoint.sh / RUN chmod +x /entrypoint.sh ENTRYPOINT ["/entrypoint.sh"] FROM base AS gitlab_webhook -ADD pr_agent pr_agent +COPY pr_agent pr_agent CMD ["python", "-m", "gunicorn", "-k", "uvicorn.workers.UvicornWorker", "-c", "pr_agent/servers/gunicorn_config.py", "--forwarded-allow-ips", "*", "pr_agent.servers.gitlab_webhook:app"] FROM base AS azure_devops_webhook -ADD pr_agent pr_agent +COPY pr_agent pr_agent CMD ["python", "pr_agent/servers/azuredevops_server_webhook.py"] FROM base AS gitea_app -ADD pr_agent pr_agent +COPY pr_agent pr_agent CMD ["python", "-m", "gunicorn", "-k", "uvicorn.workers.UvicornWorker", "-c", "pr_agent/servers/gunicorn_config.py","pr_agent.servers.gitea_app:app"] FROM base AS mosaico_agent -ADD pr_agent pr_agent +COPY pr_agent pr_agent EXPOSE 9000 CMD ["python", "pr_agent/mosaico/server.py"] FROM base AS test RUN uv sync --frozen --no-install-project -ADD pr_agent pr_agent -ADD tests tests -ADD .github/workflows/publish.yml .github/workflows/publish.yml -ADD docker/mosaico docker/mosaico +COPY pr_agent pr_agent +COPY tests tests +COPY .github/workflows/publish.yml .github/workflows/publish.yml +COPY docker/mosaico docker/mosaico FROM base AS cli -ADD pr_agent pr_agent +COPY pr_agent pr_agent RUN uv sync --frozen --no-dev && ln -s /app/.venv/bin/pr-agent /usr/local/bin/pr-agent ENTRYPOINT ["python", "pr_agent/cli.py"] diff --git a/docker/Dockerfile.lambda b/docker/Dockerfile.lambda index bda5ee05ac..112c2de12b 100644 --- a/docker/Dockerfile.lambda +++ b/docker/Dockerfile.lambda @@ -8,7 +8,8 @@ RUN dnf update -y && \ # Build the locked environment at /opt/venv with uv; put its site-packages # on the Lambda runtime's import path alongside the task root via PYTHONPATH. ENV UV_COMPILE_BYTECODE=1 UV_PROJECT_ENVIRONMENT=/opt/venv -ADD pyproject.toml uv.lock ./ +WORKDIR ${LAMBDA_TASK_ROOT} +COPY pyproject.toml uv.lock ./ RUN uv sync --frozen --no-install-project --no-dev --extra lambda --python /var/lang/bin/python3.12 ENV PYTHONPATH="/opt/venv/lib/python3.12/site-packages:${LAMBDA_TASK_ROOT}" COPY pr_agent/ ${LAMBDA_TASK_ROOT}/pr_agent/ diff --git a/docker/mosaico/README.md b/docker/mosaico/README.md index 99f7be1884..47234847d2 100644 --- a/docker/mosaico/README.md +++ b/docker/mosaico/README.md @@ -4,13 +4,13 @@ Deployment assets for running **PR-Agent** as a [MOSAICO](https://mosaico-projec *solution agent*. This directory contains no Python and no pr-agent source — it consumes PR-Agent as a published, version-pinned Docker image. The agent's source lives in [`The-PR-Agent/pr-agent`](https://github.com/The-PR-Agent/pr-agent), under `pr_agent/mosaico/`; -it is merged into `main` and ships in every release wheel and image starting at `v0.37.0`. +it is merged into `main` and ships in every release wheel and image starting at `v0.36.0`. ## Relationship to the PR-Agent repository This is not a fork and it never becomes one: -- The MOSAICO A2A server is PR-Agent code, released in tags `v0.37.0` onwards. This bundle +- The MOSAICO A2A server is PR-Agent code, released in tags `v0.36.0` onwards. This bundle holds zero Python — only a compose overlay, a registration template, an env template, a smoke test, and this README. - **Staying current is one line**: bump the pinned tag in `docker-compose.pr-agent.yml`, then diff --git a/docs/docs/assets/Ask_line.png b/docs/docs/assets/Ask_line.png new file mode 100644 index 0000000000..848132c349 Binary files /dev/null and b/docs/docs/assets/Ask_line.png differ diff --git a/docs/docs/assets/add_docs_comment.png b/docs/docs/assets/add_docs_comment.png new file mode 100644 index 0000000000..5ff47a8cbb Binary files /dev/null and b/docs/docs/assets/add_docs_comment.png differ diff --git a/docs/docs/assets/add_native_custom_labels.webp b/docs/docs/assets/add_native_custom_labels.webp new file mode 100644 index 0000000000..79b25f3500 Binary files /dev/null and b/docs/docs/assets/add_native_custom_labels.webp differ diff --git a/docs/docs/assets/ask.png b/docs/docs/assets/ask.png new file mode 100644 index 0000000000..e37705e189 Binary files /dev/null and b/docs/docs/assets/ask.png differ diff --git a/docs/docs/assets/ask_comment.png b/docs/docs/assets/ask_comment.png new file mode 100644 index 0000000000..cdb5d5645b Binary files /dev/null and b/docs/docs/assets/ask_comment.png differ diff --git a/docs/docs/assets/ask_images1.png b/docs/docs/assets/ask_images1.png new file mode 100644 index 0000000000..06f401ccf1 Binary files /dev/null and b/docs/docs/assets/ask_images1.png differ diff --git a/docs/docs/assets/ask_images2.png b/docs/docs/assets/ask_images2.png new file mode 100644 index 0000000000..15dae76472 Binary files /dev/null and b/docs/docs/assets/ask_images2.png differ diff --git a/docs/docs/assets/ask_images3.png b/docs/docs/assets/ask_images3.png new file mode 100644 index 0000000000..d788afe3a3 Binary files /dev/null and b/docs/docs/assets/ask_images3.png differ diff --git a/docs/docs/assets/ask_images4.png b/docs/docs/assets/ask_images4.png new file mode 100644 index 0000000000..63cdc96b38 Binary files /dev/null and b/docs/docs/assets/ask_images4.png differ diff --git a/docs/docs/assets/ask_images5.png b/docs/docs/assets/ask_images5.png new file mode 100644 index 0000000000..0ea9fb7a24 Binary files /dev/null and b/docs/docs/assets/ask_images5.png differ diff --git a/docs/docs/assets/azure_devops_pipeline_secrets.png b/docs/docs/assets/azure_devops_pipeline_secrets.png new file mode 100644 index 0000000000..aa838b3e41 Binary files /dev/null and b/docs/docs/assets/azure_devops_pipeline_secrets.png differ diff --git a/docs/docs/assets/code_suggestions_as_comment_closed.png b/docs/docs/assets/code_suggestions_as_comment_closed.png new file mode 100644 index 0000000000..74d928f9eb Binary files /dev/null and b/docs/docs/assets/code_suggestions_as_comment_closed.png differ diff --git a/docs/docs/assets/code_suggestions_as_comment_open.png b/docs/docs/assets/code_suggestions_as_comment_open.png new file mode 100644 index 0000000000..08bad76a8b Binary files /dev/null and b/docs/docs/assets/code_suggestions_as_comment_open.png differ diff --git a/docs/docs/assets/describe_comment.png b/docs/docs/assets/describe_comment.png new file mode 100644 index 0000000000..5b13930ddb Binary files /dev/null and b/docs/docs/assets/describe_comment.png differ diff --git a/docs/docs/assets/describe_markers_after.webp b/docs/docs/assets/describe_markers_after.webp new file mode 100644 index 0000000000..fe021244a5 Binary files /dev/null and b/docs/docs/assets/describe_markers_after.webp differ diff --git a/docs/docs/assets/describe_markers_before.png b/docs/docs/assets/describe_markers_before.png new file mode 100644 index 0000000000..939099c80d Binary files /dev/null and b/docs/docs/assets/describe_markers_before.png differ diff --git a/docs/docs/assets/describe_new.webp b/docs/docs/assets/describe_new.webp new file mode 100644 index 0000000000..8fc0b222d5 Binary files /dev/null and b/docs/docs/assets/describe_new.webp differ diff --git a/docs/docs/assets/describe_new_short_main.webp b/docs/docs/assets/describe_new_short_main.webp new file mode 100644 index 0000000000..26cfca0f2e Binary files /dev/null and b/docs/docs/assets/describe_new_short_main.webp differ diff --git a/docs/docs/assets/diagram-v0.9.webp b/docs/docs/assets/diagram-v0.9.webp new file mode 100644 index 0000000000..1c0694cf3c Binary files /dev/null and b/docs/docs/assets/diagram-v0.9.webp differ diff --git a/docs/docs/assets/filter_mail_notifications.png b/docs/docs/assets/filter_mail_notifications.png new file mode 100644 index 0000000000..329b8efa8f Binary files /dev/null and b/docs/docs/assets/filter_mail_notifications.png differ diff --git a/docs/docs/assets/git_patch_logic.png b/docs/docs/assets/git_patch_logic.png new file mode 100644 index 0000000000..b88c382ab0 Binary files /dev/null and b/docs/docs/assets/git_patch_logic.png differ diff --git a/docs/docs/assets/help1.png b/docs/docs/assets/help1.png new file mode 100644 index 0000000000..04ffae16c4 Binary files /dev/null and b/docs/docs/assets/help1.png differ diff --git a/docs/docs/assets/help2.webp b/docs/docs/assets/help2.webp new file mode 100644 index 0000000000..9519e72d09 Binary files /dev/null and b/docs/docs/assets/help2.webp differ diff --git a/docs/docs/assets/help_docs_comment.png b/docs/docs/assets/help_docs_comment.png new file mode 100644 index 0000000000..a63014a438 Binary files /dev/null and b/docs/docs/assets/help_docs_comment.png differ diff --git a/docs/docs/assets/help_docs_comment_explicit_git.png b/docs/docs/assets/help_docs_comment_explicit_git.png new file mode 100644 index 0000000000..f1ef78dd12 Binary files /dev/null and b/docs/docs/assets/help_docs_comment_explicit_git.png differ diff --git a/docs/docs/assets/help_docs_response.png b/docs/docs/assets/help_docs_response.png new file mode 100644 index 0000000000..9200d3f4d2 Binary files /dev/null and b/docs/docs/assets/help_docs_response.png differ diff --git a/docs/docs/assets/improve.png b/docs/docs/assets/improve.png new file mode 100644 index 0000000000..04b413dcd0 Binary files /dev/null and b/docs/docs/assets/improve.png differ diff --git a/docs/docs/assets/improve_new_short_main.webp b/docs/docs/assets/improve_new_short_main.webp new file mode 100644 index 0000000000..b6e5a82161 Binary files /dev/null and b/docs/docs/assets/improve_new_short_main.webp differ diff --git a/docs/docs/assets/notifications.png b/docs/docs/assets/notifications.png new file mode 100644 index 0000000000..f8ce3089f6 Binary files /dev/null and b/docs/docs/assets/notifications.png differ diff --git a/docs/docs/assets/possible_config1.png b/docs/docs/assets/possible_config1.png new file mode 100644 index 0000000000..e09246055d Binary files /dev/null and b/docs/docs/assets/possible_config1.png differ diff --git a/docs/docs/assets/possible_config2.png b/docs/docs/assets/possible_config2.png new file mode 100644 index 0000000000..9d123283b7 Binary files /dev/null and b/docs/docs/assets/possible_config2.png differ diff --git a/docs/docs/assets/pr_description_user_description.png b/docs/docs/assets/pr_description_user_description.png new file mode 100644 index 0000000000..8665fc3a82 Binary files /dev/null and b/docs/docs/assets/pr_description_user_description.png differ diff --git a/docs/docs/assets/review3.png b/docs/docs/assets/review3.png new file mode 100644 index 0000000000..3fd9e89894 Binary files /dev/null and b/docs/docs/assets/review3.png differ diff --git a/docs/docs/assets/review_comment.png b/docs/docs/assets/review_comment.png new file mode 100644 index 0000000000..17dc535f00 Binary files /dev/null and b/docs/docs/assets/review_comment.png differ diff --git a/docs/docs/assets/review_new_short_main.png b/docs/docs/assets/review_new_short_main.png new file mode 100644 index 0000000000..13c86c915a Binary files /dev/null and b/docs/docs/assets/review_new_short_main.png differ diff --git a/docs/docs/assets/self_reflection1.png b/docs/docs/assets/self_reflection1.png new file mode 100644 index 0000000000..56c8bfda2a Binary files /dev/null and b/docs/docs/assets/self_reflection1.png differ diff --git a/docs/docs/assets/self_reflection2.png b/docs/docs/assets/self_reflection2.png new file mode 100644 index 0000000000..e9bbb219eb Binary files /dev/null and b/docs/docs/assets/self_reflection2.png differ diff --git a/docs/docs/assets/self_review_1.webp b/docs/docs/assets/self_review_1.webp new file mode 100644 index 0000000000..c7404164c4 Binary files /dev/null and b/docs/docs/assets/self_review_1.webp differ diff --git a/docs/docs/assets/similar_issue.png b/docs/docs/assets/similar_issue.png new file mode 100644 index 0000000000..5f5d995c08 Binary files /dev/null and b/docs/docs/assets/similar_issue.png differ diff --git a/docs/docs/assets/similar_issue_comment.png b/docs/docs/assets/similar_issue_comment.png new file mode 100644 index 0000000000..0e8d2efc8f Binary files /dev/null and b/docs/docs/assets/similar_issue_comment.png differ diff --git a/docs/docs/assets/similar_issue_original_issue.png b/docs/docs/assets/similar_issue_original_issue.png new file mode 100644 index 0000000000..a09e703aab Binary files /dev/null and b/docs/docs/assets/similar_issue_original_issue.png differ diff --git a/docs/docs/assets/ticket_compliance_review.png b/docs/docs/assets/ticket_compliance_review.png new file mode 100644 index 0000000000..46e883ea79 Binary files /dev/null and b/docs/docs/assets/ticket_compliance_review.png differ diff --git a/docs/docs/assets/update_changelog.png b/docs/docs/assets/update_changelog.png new file mode 100644 index 0000000000..eb2472a4cc Binary files /dev/null and b/docs/docs/assets/update_changelog.png differ diff --git a/docs/docs/assets/update_changelog_comment.png b/docs/docs/assets/update_changelog_comment.png new file mode 100644 index 0000000000..deeb8c27aa Binary files /dev/null and b/docs/docs/assets/update_changelog_comment.png differ diff --git a/docs/docs/core-abilities/compression_strategy.md b/docs/docs/core-abilities/compression_strategy.md index a4448f0e0b..a7983bb497 100644 --- a/docs/docs/core-abilities/compression_strategy.md +++ b/docs/docs/core-abilities/compression_strategy.md @@ -54,4 +54,4 @@ We use [tiktoken](https://github.com/openai/tiktoken) to tokenize the patches af #### Example -![Core Abilities](https://codium.ai/images/git_patch_logic.png){width=768} +![Core Abilities](../assets/git_patch_logic.png){width=768} diff --git a/docs/docs/core-abilities/fetching_ticket_context.md b/docs/docs/core-abilities/fetching_ticket_context.md index 7ef9cf910c..64e5c64445 100644 --- a/docs/docs/core-abilities/fetching_ticket_context.md +++ b/docs/docs/core-abilities/fetching_ticket_context.md @@ -1,9 +1,10 @@ # Fetching Ticket Context for PRs -`Supported Git Platforms: GitHub, GitLab, Bitbucket` +`Supported Git Platforms: GitHub, GitLab, Bitbucket, Azure DevOps` -!!! note "Branch-name issue linking: GitHub only (for now)" - Extracting issue links from the **branch name** (and the optional `branch_issue_regex` setting) is currently implemented for **GitHub only**. Support for GitLab, Bitbucket, and other platforms is planned for a later release. The GitHub flow was the most relevant to implement first; other providers will follow. +!!! note "Branch-name linking: Jira keys on all providers; numeric GitHub issues on GitHub only" + **Jira** ticket keys (e.g. `ABC-123`) are extracted from the branch name on **every git provider**. + Extracting **numeric GitHub issue** links from the branch name (and the optional `branch_issue_regex` setting) is currently implemented for **GitHub only**; support for other providers is planned for a later release. ## Overview @@ -50,7 +51,7 @@ Each ticket will be assigned a label (Compliance/Alignment level), Indicates the - Not Compliant - PR Code Verified -![Ticket Compliance](https://www.qodo.ai/images/pr_agent/ticket_compliance_review.png){width=768} +![Ticket Compliance](../assets/ticket_compliance_review.png){width=768} A `PR Code Verified` label indicates the PR code meets ticket requirements, but requires additional manual testing beyond the code scope. For example - validating UI display across different environments (Mac, Windows, mobile, etc.). @@ -147,7 +148,11 @@ the established ticket-extraction behavior of existing providers. ## Jira Integration -We support both Jira Cloud and Jira Server/Data Center. +Only **Jira Cloud** is supported. The base URL is derived from a validated site name +(`jira_site` → `https://.atlassian.net`) rather than taken as a free-form URL, so +the configured destination is always an Atlassian Cloud host. Jira Server / Data Center +(self-hosted) uses a free-form host and is not supported yet; it can be added once +base-URL handling for the self-hosted case is settled. ### Jira Cloud @@ -169,187 +174,48 @@ You can create an API token from your Atlassian account: ```toml [jira] -jira_api_token = "YOUR_API_TOKEN" +jira_site = "" # the "" in https://.atlassian.net (e.g. "mycompany") jira_api_email = "YOUR_EMAIL" +jira_api_token = "YOUR_API_TOKEN" ``` -### Jira Data Center/Server +`jira_site` is your Jira Cloud site name — the part before `.atlassian.net` (for +`https://mycompany.atlassian.net`, the site is `mycompany`). PR-Agent builds the base URL +as `https://.atlassian.net`; it does not accept a full URL, so configuration +cannot redirect the authenticated request to another host. Store `jira_api_email` and +`jira_api_token` as secrets (environment variables or the secrets file), not in +repository-committed configuration. -#### Using Basic Authentication for Jira Data Center/Server +#### Acceptance criteria / requirements (optional) -You can use your Jira username and password to authenticate with Jira Data Center/Server. - -In your Configuration file/Environment variables/Secrets file, add the following lines: +To include a ticket's acceptance criteria in the analysis, set `jira_requirements_field` +to the id of the custom field that holds it. The field id is specific to your Jira +instance (for example `customfield_10127`); leave it empty to skip requirements. ```toml -jira_api_email = "your_username" -jira_api_token = "your_password" +[jira] +jira_requirements_field = "customfield_10127" ``` -(Note that indeed the 'jira_api_email' field is used for the username, and the 'jira_api_token' field is used for the user password.) - -##### Validating Basic authentication via Python script - -If you are facing issues retrieving tickets in PR-Agent with Basic auth, you can validate the flow using a Python script. -This following steps will help you check if the basic auth is working correctly, and if you can access the Jira ticket details: - -1. run `pip install jira==3.8.0` - -2. run the following Python script (after replacing the placeholders with your actual values): - -???- example "Script to validate basic auth" - - ```python - from jira import JIRA - - - if __name__ == "__main__": - try: - # Jira server URL - server = "https://..." - # Basic auth - username = "..." - password = "..." - # Jira ticket code (e.g. "PROJ-123") - ticket_id = "..." - - print("Initializing JiraServerTicketProvider with JIRA server") - # Initialize JIRA client - jira = JIRA( - server=server, - basic_auth=(username, password), - timeout=30 - ) - if jira: - print(f"JIRA client initialized successfully") - else: - print("Error initializing JIRA client") - - # Fetch ticket details - ticket = jira.issue(ticket_id) - print(f"Ticket title: {ticket.fields.summary}") - - except Exception as e: - print(f"Error fetching JIRA ticket details: {e}") - ``` - -#### Using a Personal Access Token (PAT) for Jira Data Center/Server +#### Project key allowlist (optional) -1. Create a [Personal Access Token (PAT)](https://confluence.atlassian.com/enterprise/using-personal-access-tokens-1026032365.html) in your Jira account -2. In your Configuration file/Environment variables/Secrets file, add the following lines: +Ticket detection matches any `PROJECT-123` shaped text, so strings like `SHA-256`, +`UTF-8` or `ISO-8601` in a title or description each cost an authenticated lookup that +returns 404. If your repository works with a known set of Jira projects, list their keys +in `project_keys`; keys with any other prefix are then dropped before any lookup (they are +named once at debug level in the log). Leave the list empty to look up every key found. ```toml [jira] -jira_base_url = "YOUR_JIRA_BASE_URL" # e.g. https://jira.example.com -jira_api_token = "YOUR_API_TOKEN" +project_keys = ["PROJ", "OPS"] ``` -##### Validating PAT token via Python script - -If you are facing issues retrieving tickets in PR-Agent with PAT token, you can validate the flow using a Python script. -This following steps will help you check if the token is working correctly, and if you can access the Jira ticket details: - -1. run `pip install jira==3.8.0` - -2. run the following Python script (after replacing the placeholders with your actual values): - -??? example- "Script to validate PAT token" - - ```python - from jira import JIRA - - - if __name__ == "__main__": - try: - # Jira server URL - server = "https://..." - # Jira PAT token - token_auth = "..." - # Jira ticket code (e.g. "PROJ-123") - ticket_id = "..." - - print("Initializing JiraServerTicketProvider with JIRA server") - # Initialize JIRA client - jira = JIRA( - server=server, - token_auth=token_auth, - timeout=30 - ) - if jira: - print(f"JIRA client initialized successfully") - else: - print("Error initializing JIRA client") - - # Fetch ticket details - ticket = jira.issue(ticket_id) - print(f"Ticket title: {ticket.fields.summary}") - - except Exception as e: - print(f"Error fetching JIRA ticket details: {e}") - ``` - - -### Multi-JIRA Server Configuration - -PR-Agent supports connecting to multiple JIRA servers using different authentication methods. - -=== "Email/Token (Basic Auth)" - - Configure multiple servers using Email/Token authentication: - - - `jira_servers`: List of JIRA server URLs - - `jira_api_token`: List of API tokens (for Cloud) or passwords (for Data Center) - - `jira_api_email`: List of emails (for Cloud) or usernames (for Data Center) - - `jira_base_url`: Default server for ticket IDs like `PROJ-123`, Each repository can configure (local config file) its own `jira_base_url` to choose which server to use by default. - - **Example Configuration:** - ```toml - [jira] - # Server URLs - jira_servers = ["https://company.atlassian.net", "https://datacenter.jira.com"] - - # API tokens/passwords - jira_api_token = ["cloud_api_token_here", "datacenter_password"] - - # Emails/usernames (both required) - jira_api_email = ["user@company.com", "datacenter_username"] - - # Default server for ticket IDs - jira_base_url = "https://company.atlassian.net" - ``` - -=== "PAT Auth" - - Configure multiple servers using Personal Access Token authentication: - - - `jira_servers`: List of JIRA server URLs - - `jira_api_token`: List of PAT tokens - - `jira_api_email`: Not needed (can be omitted or left empty) - - `jira_base_url`: Default server for ticket IDs like `PROJ-123`, Each repository can configure (local config file) its own `jira_base_url` to choose which server to use by default. - - **Example Configuration:** - ```toml - [jira] - # Server URLs - jira_servers = ["https://server1.jira.com", "https://server2.jira.com"] - - # PAT tokens only - jira_api_token = ["pat_token_1", "pat_token_2"] - - # Default server for ticket IDs - jira_base_url = "https://server1.jira.com" - ``` - - **Mixed Authentication (Email/Token + PAT):** - ```toml - [jira] - jira_servers = ["https://company.atlassian.net", "https://server.jira.com"] - jira_api_token = ["cloud_api_token", "server_pat_token"] - jira_api_email = ["user@company.com", ""] # Empty for PAT - ``` - - - +Entries are plain upper-case project keys (letters only, as Jira writes them); anything +else (a lower-case label, a full ticket key, a URL, a blank entry) is ignored with a +warning. If the list is set but none of its entries is valid, no Jira lookup is made at all +until it is fixed, so a typo cannot silently widen the lookup again. Only a missing option, +the empty list, and an environment override set to the empty string mean "look up every +key". ### How to link a PR to a Jira ticket @@ -357,17 +223,8 @@ To integrate with Jira, you can link your PR to a ticket using either of these m **Method 1: Description Reference:** -Include a ticket reference in your PR description, using either the complete URL format `https://.atlassian.net/browse/ISSUE-123` or the shortened ticket ID `ISSUE-123` (without prefix or suffix for the shortened ID). +Include a ticket reference in your PR description, using either the complete URL format `https://.atlassian.net/browse/ISSUE-123` or the shortened ticket ID `ISSUE-123` (without prefix or suffix for the shortened ID). **Method 2: Branch Name Detection:** Name your branch with the ticket ID as a prefix (e.g., `ISSUE-123-feature-description` or `ISSUE-123/feature-description`). - -!!! note "Jira Base URL" - For shortened ticket IDs or branch detection (method 2 for JIRA cloud), you must configure the Jira base URL in your configuration file under the [jira] section: - - ```toml - [jira] - jira_base_url = "https://.atlassian.net" - ``` - Where `` is your Jira organization identifier (e.g., `mycompany` for `https://mycompany.atlassian.net`). diff --git a/docs/docs/core-abilities/self_reflection.md b/docs/docs/core-abilities/self_reflection.md index 990f724fbf..772afac2da 100644 --- a/docs/docs/core-abilities/self_reflection.md +++ b/docs/docs/core-abilities/self_reflection.md @@ -38,8 +38,8 @@ This results in a more refined and valuable set of suggestions for the user, sav ## Example Results -![self_reflection](https://codium.ai/images/pr_agent/self_reflection1.png){width=768} -![self_reflection](https://codium.ai/images/pr_agent/self_reflection2.png){width=768} +![self_reflection](../assets/self_reflection1.png){width=768} +![self_reflection](../assets/self_reflection2.png){width=768} ## Appendix - Relevant Configuration Options diff --git a/docs/docs/index.md b/docs/docs/index.md index 6e4510c794..9ac27c6304 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -58,21 +58,21 @@ PR-Agent offers comprehensive pull request functionalities integrated with vario #### [/describe](https://github.com/the-pr-agent/pr-agent/pull/530)
-![/describe](https://www.codium.ai/images/pr_agent/describe_new_short_main.png){width=512} +![/describe](assets/describe_new_short_main.webp){width=512}
#### [/review](https://github.com/the-pr-agent/pr-agent/pull/732#issuecomment-1975099151)
-![/review](https://www.codium.ai/images/pr_agent/review_new_short_main.png){width=512} +![/review](assets/review_new_short_main.png){width=512}
#### [/improve](https://github.com/the-pr-agent/pr-agent/pull/732#issuecomment-1975099159)
-![/improve](https://www.codium.ai/images/pr_agent/improve_new_short_main.png){width=512} +![/improve](assets/improve_new_short_main.webp){width=512}
@@ -80,6 +80,6 @@ PR-Agent offers comprehensive pull request functionalities integrated with vario The following diagram illustrates PR-Agent tools and their flow: -![PR-Agent Tools](https://codium.ai/images/pr_agent/diagram-v0.9.png) +![PR-Agent Tools](assets/diagram-v0.9.webp) Check out the [PR Compression strategy](core-abilities/index.md) page for more details on how we convert a code diff to a manageable LLM prompt diff --git a/docs/docs/installation/azure.md b/docs/docs/installation/azure.md index 817d67f669..61a61d93d9 100644 --- a/docs/docs/installation/azure.md +++ b/docs/docs/installation/azure.md @@ -59,7 +59,7 @@ stages: This script will run PR-Agent on every new merge request, with the `improve`, `review`, and `describe` commands. Note that you need to export the `azure_devops__pat` and `OPENAI_KEY` variables in the Azure DevOps pipeline settings (Pipelines -> Library -> + Variable group): -![PR-Agent](https://codium.ai/images/pr_agent/azure_devops_pipeline_secrets.png){width=468} +![PR-Agent](../assets/azure_devops_pipeline_secrets.png){width=468} Make sure to give pipeline permissions to the `pr_agent` variable group. diff --git a/docs/docs/installation/github.md b/docs/docs/installation/github.md index 9e3da8a862..ef708a040e 100644 --- a/docs/docs/installation/github.md +++ b/docs/docs/installation/github.md @@ -616,7 +616,7 @@ For more detailed configuration options, see: ... ``` - For enhanced security, you can also specify the Docker image by its [digest](https://hub.docker.com/repository/docker/pragent/pr-agent/tags). Resolve the digest for the version you are pinning with `docker buildx imagetools inspect pragent/pr-agent:0.41.0-github_action --format '{{.Manifest.Digest}}'`, then use it in place of the tag: + For enhanced security, you can also specify the Docker image by its [digest](https://hub.docker.com/r/pragent/pr-agent/tags). Resolve the digest for the version you are pinning with `docker buildx imagetools inspect pragent/pr-agent:0.41.0-github_action --format '{{.Manifest.Digest}}'`, then use it in place of the tag: ```yaml ... steps: @@ -845,6 +845,7 @@ Example IAM permissions to that user to allow access to CodeCommit: "codecommit:List*", "codecommit:PostComment*", "codecommit:PutCommentReaction", + "codecommit:UpdateComment", "codecommit:UpdatePullRequestDescription", "codecommit:UpdatePullRequestTitle" ], diff --git a/docs/docs/installation/gitlab.md b/docs/docs/installation/gitlab.md index d4ccbdae2d..447bfb5687 100644 --- a/docs/docs/installation/gitlab.md +++ b/docs/docs/installation/gitlab.md @@ -88,6 +88,7 @@ Note that if your base branches are not protected, don't set the variables as `p > **Note**: The `gitlab__SSL_VERIFY` environment variable can be used to specify the path to a custom CA certificate bundle for SSL verification. GitLab exposes the `$CI_SERVER_TLS_CA_FILE` variable, which points to the custom CA certificate file configured in your GitLab instance. > Alternatively, SSL verification can be disabled entirely by setting `gitlab__SSL_VERIFY=false`, although this is not recommended. +> This setting affects only the GitLab API client. For certificate issues with LLM calls or git clone operations, see [Custom CA and Self-Signed Certificates](../usage-guide/custom_ca_and_self_signed_certificates.md). ## Run a GitLab webhook server diff --git a/docs/docs/installation/index.md b/docs/docs/installation/index.md index e008d8a7aa..a8fa601ab4 100644 --- a/docs/docs/installation/index.md +++ b/docs/docs/installation/index.md @@ -23,7 +23,7 @@ behavior, not Requests-specific environment settings. ## Sizing a self-hosted webhook server -The GitHub, GitLab and Gitea webhook servers (the `github_app`, `gitlab_webhook` and `gitea_app` Docker targets) run under gunicorn with multiple worker processes, so that a worker busy handling a request cannot block the health check served by another. The other deployments — Bitbucket, Azure DevOps, GitHub polling, and the Lambda variants — run a single process and are unaffected by this section. +The GitHub, GitLab and Gitea webhook servers (the `github_app`, `gitlab_webhook` and `gitea_app` Docker targets) run under gunicorn with multiple worker processes, so that a worker busy handling a request cannot block the health check served by another. The other deployments — Bitbucket, Azure DevOps, GitHub polling, and the Lambda variants — do not use this gunicorn worker configuration. | Variable | Default | Description | |----------|---------|-------------| @@ -40,6 +40,24 @@ If the container is OOMKilled during startup, lower the worker count: GUNICORN_WORKERS=2 ``` +## GitHub polling workers + +GitHub polling uses one polling process and a separate child process for each +accepted comment. At most 10 of these child workers may be active across all +polling iterations. If an earlier batch still fills the limit, accepted work +waits for a free slot without blocking the event loop. This can delay the next +notification poll under sustained load. + +The existing per-batch limit is unchanged: only the first 10 queued comments in +a batch are accepted; excess work is logged and dropped. Waiting work is held in +memory only. Notification acknowledgement is unchanged, so shutdown or startup +failures do not guarantee redelivery. Dispatch interruptions log the remaining +count. A worker startup failure stops polling rather than risking further +dispatch with an untracked child; inspect the failure before restarting. +Completed children are joined without waiting and closed each iteration; +live children are not terminated on cancellation. This is not a durable queue, +a model-call deadline, or a host-wide memory limit. + !!! note "Docker Hub namespace migration" Releases **`0.34.2` and later** are published under [`pragent/pr-agent`](https://hub.docker.com/r/pragent/pr-agent). Older releases (up to and including `v0.31`) remain at the legacy [`codiumai/pr-agent`](https://hub.docker.com/r/codiumai/pr-agent) namespace as a frozen archive — no new images are pushed there. The examples on this site reference the new namespace; if you are pinning to a release before `0.34.2`, swap `pragent/pr-agent` for `codiumai/pr-agent` in your `image:` / `docker pull` / `uses: docker://` references. diff --git a/docs/docs/installation/locally.md b/docs/docs/installation/locally.md index 3dc7da908e..c4fb9e6f2f 100644 --- a/docs/docs/installation/locally.md +++ b/docs/docs/installation/locally.md @@ -122,13 +122,15 @@ def main(): get_settings().set("github.user_token", user_token) # Run the command. Feedback will appear in GitHub PR comments - cli.run_command(pr_url, command) + return cli.run_command(pr_url, command) if __name__ == '__main__': - main() + raise SystemExit(main()) ``` +With `config.propagate_tool_errors` enabled, forwarding the return value through `SystemExit` makes this script exit with status 1 after a propagated tool error. The default remains status 0. + ## Run from source 1. Clone this repository: diff --git a/docs/docs/installation/mosaico_server.md b/docs/docs/installation/mosaico_server.md new file mode 100644 index 0000000000..f4e3b04275 --- /dev/null +++ b/docs/docs/installation/mosaico_server.md @@ -0,0 +1,132 @@ +## MOSAICO A2A server + +PR-Agent can run as an [A2A](https://a2a-protocol.org/) 1.0 *solution agent* for the +[MOSAICO](https://mosaico-project.eu/) ecosystem: a small Starlette server that exposes the +standard A2A surface (agent card + JSON-RPC) plus a health probe. It is **not** a fork or a +separate project — the server is PR-Agent code under [`pr_agent/mosaico/`](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/mosaico/server.py), +ships in every release wheel, and ships as its own Docker image (`-mosaico_agent`) +starting at `v0.36.0`. The server is unbiased about the git provider: every request carries +either a PR URL or a raw diff, and the agent answers from that input alone. + +### What the mode is + +The A2A server exposes three endpoints: + +| Path | Method | Purpose | +| --- | --- | --- | +| `/.well-known/agent-card.json` | GET | The A2A agent card | +| `/` | POST | A2A 1.0 JSON-RPC. Sends the `SendMessage` method with an `A2A-Version: 1.0` header (the header is required: the server treats a request without it as protocol 0.3 and rejects it). The reply arrives as a task artifact (`result.task.artifacts[].parts[].text`), not as a status message. | +| `/health` | GET | A **live LLM connectivity probe** — `200` when an LLM round-trip succeeds, `503` otherwise | + +The advertised agent card carries the skills `review`, `improve`, `describe`, and `ask`, the +name `"PR-Agent Solution Agent"`, a `version` derived from the running build (never +hand-maintained), and the required +`https://mosaico-project.eu/extensions/mosaico-observability` extension. Streaming is +advertised as disabled, which is load-bearing: the reference agent selects +`message/send` vs `message/stream` from that capability. + +### Run the standalone container + +The server boots from a bare `docker pull` in a couple of seconds — no repo clone, no build: + +```bash +docker pull pragent/pr-agent:0.41.0-mosaico_agent +docker run -d --name pr-agent-mosaico -p 9000:9000 \ + -e API_BASE=https://your-openai-compatible-endpoint/v1 \ + -e API_KEY=sk-... \ + -e MODEL_NAME=openai/your-model-slug \ + pragent/pr-agent:0.41.0-mosaico_agent + +curl -s http://localhost:9000/.well-known/agent-card.json | python3 -m json.tool +``` + +Pin a version tag in production (see the "Immutable releases and version tags" note on the +[installation page](./index.md)); the plain `mosaico_agent` +rolling tag moves to the newest build on every release. + +### Environment variables + +| Variable | Default | Purpose | +| --- | --- | --- | +| `API_BASE` | — | Base URL of the OpenAI-compatible LLM endpoint | +| `API_KEY` | — | API key for that endpoint | +| `MODEL_NAME` | — | Model slug to call | +| `HOST` | `0.0.0.0` | Bind address | +| `PORT` | `9000` | Bind port | +| `AGENT_CARD_HOST`, `AGENT_CARD_PORT` | unset | URL advertised in the card's `supportedInterfaces`; see the warning below | +| `MODEL_MAX_TOKENS` | `32000` | Token budget for models whose context size pr-agent does not already know | +| `LANGFUSE_HOST`, `LANGFUSE_PUBLIC_KEY`, `LANGFUSE_SECRET_KEY` | unset | Optional Langfuse observability | + +!!! warning "AGENT_CARD_HOST / AGENT_CARD_PORT — the one thing to get right" + + These two variables set the URL the agent advertises in `supportedInterfaces`. Leave them + unset and the card advertises `http://localhost:9000/`, which is reachable only from + inside the container itself. The failure this causes is **silent and late**: registration + with MOSAICO succeeds, the repository stores the unreachable URL, and the reference agent + only fails to dereference it once it tries to route a task to this agent. Set them to + whatever host/port the *caller* will use to reach the container, and verify with: + + ```bash + curl -s http://:/.well-known/agent-card.json \ + | python3 -c "import sys,json; print(json.load(sys.stdin)['supportedInterfaces'][0]['url'])" + ``` + + If that prints a `localhost` URL, the deployment is wrong. + +### Deploy into the mosaico-demonstrator + +The [`docker/mosaico/`](https://github.com/the-pr-agent/pr-agent/tree/main/docker/mosaico) +directory is a full deployment bundle (compose overlay, registration template, env template, +smoke test, LICENSE, and the canonical README). To run the agent as a task agent in the +[mosaico-demonstrator](https://gitlab.eclipse.org/eclipse-research-labs/mosaico-project/mosaico-demonstrator): + +1. Copy `docker-compose.pr-agent.yml` into the demonstrator's `compose/` directory, next to + `base-definitions.yml` (the overlay's `extends:` references resolve relative to that + directory). +2. Copy `pr-agent-solution-agent.json` into the demonstrator's + `docker/agent-registrations/` directory. +3. Append the "demonstrator overlay" block from `pr-agent.env.example` to the demonstrator's + `env/llm.env` and fill in `PR_AGENT_MODEL` (`PR_AGENT_HOST` may stay empty to use the + demonstrator's auto-detected LAN IP; `PR_AGENT_PORT` defaults to `23000`). +4. Add `-f compose/docker-compose.pr-agent.yml` to the demonstrator's `01-compose.sh`, next to + the other task-agent overlays. +5. Run `./01-compose.sh up -d`. + +The registration template carries only `description`, `role`, `objective`, `version`; the +demonstrator's `register-agent.py` injects `name`, `a2aAgentCardUrl`, and +`deployment.mode = ENDPOINT` at registration time. Two names are intentionally different and +should not be "fixed": the repository entry is `pr-agent-solution-agent` (what +`register-agent.py` looks the agent up by), while the card's own `name` is +`"PR-Agent Solution Agent"` (a display string). + +### Verify + +```bash +./smoke_test.sh +``` + +in the bundle directory gives one of two outcomes: + +- **`SMOKE PASSED`** — no LLM credentials were available; the script pulled the pinned image, + booted it, and validated the agent card only. +- **`FULL ROUND-TRIP PASSED`** — credentials were present (via a `.env` file beside the script, + copied from `pr-agent.env.example`); the script additionally exercised `GET /health` and an + A2A `SendMessage` review over an inline diff. + +### Troubleshooting + +- **The container stays `unhealthy` and registration never runs.** `/health` is a live LLM + probe and returns `503` on bad or missing credentials — this is intended. Check + `API_BASE` / `API_KEY` / `MODEL_NAME`, not the compose file. +- **The agent registers but the reference agent never reaches it.** The advertised card URL is + `localhost`; see the `AGENT_CARD_HOST` / `AGENT_CARD_PORT` warning above. +- **The registration container itself cannot fetch the agent card.** `01-compose.sh` falls back + to `get_fallback_ip`, which can resolve to `localhost` — reachable from the host but not from + inside the registration container on the Docker network. Set `PR_AGENT_HOST` explicitly to an + address reachable from inside Docker (for example the host's LAN IP, or `host.docker.internal`). + +### Keep reading + +The [bundle README](https://github.com/the-pr-agent/pr-agent/blob/main/docker/mosaico/README.md) +is the canonical deep dive for this surface and the source of the summary above; it covers the +upgrade procedure, the registration flow, and the full env-var contract in one place. diff --git a/docs/docs/summary.md b/docs/docs/summary.md index b4e71bb2f8..568c943baf 100644 --- a/docs/docs/summary.md +++ b/docs/docs/summary.md @@ -13,17 +13,21 @@ * [BitBucket Integration](installation/bitbucket.md) * [Azure DevOps Integration](installation/azure.md) * [Gitea Integration](installation/gitea.md) +* [MOSAICO A2A Server](installation/mosaico_server.md) ## Usage Guide * [Usage Guide](usage-guide/index.md) * [Introduction](usage-guide/introduction.md) * [Configuration File](usage-guide/configuration_options.md) +* [Configuration Reference](usage-guide/configuration_reference.md) * [Usage and Automation](usage-guide/automations_and_usage.md) * [Managing Mail Notifications](usage-guide/mail_notifications.md) +* [Push Outputs](usage-guide/push_outputs.md) * [Changing a Model](usage-guide/changing_a_model.md) * [Extending PR-Agent](usage-guide/extending_pr_agent.md) * [Additional Configurations](usage-guide/additional_configurations.md) +* [Custom CA and Self-Signed Certificates](usage-guide/custom_ca_and_self_signed_certificates.md) * [Plain-Diff Mode](usage-guide/plain_diff_mode.md) * [Local Git Provider](usage-guide/local_git_provider.md) * [Frequently Asked Questions](faq/index.md) diff --git a/docs/docs/tools/add_docs.md b/docs/docs/tools/add_docs.md index 02b5f6a37e..d5b926e484 100644 --- a/docs/docs/tools/add_docs.md +++ b/docs/docs/tools/add_docs.md @@ -12,11 +12,9 @@ It can be invoked manually by commenting on any PR: Invoke the tool manually by commenting `/add_docs` on any PR: -![Add Docs](https://codium.ai/images/pr_agent/add_docs_comment.png){width=512} +![Add Docs](../assets/add_docs_comment.png){width=512} -The tool will generate documentation suggestions as inline code suggestions: - -![Add Docs Result](https://codium.ai/images/pr_agent/add_docs_result.png){width=512} +The tool will generate documentation suggestions as inline code suggestions. ### Language-specific documentation styles diff --git a/docs/docs/tools/ask.md b/docs/docs/tools/ask.md index f87c982055..025d5a1d6e 100644 --- a/docs/docs/tools/ask.md +++ b/docs/docs/tools/ask.md @@ -9,9 +9,9 @@ It can be invoked manually by commenting on any PR: ## Example usage -![Ask Comment](https://codium.ai/images/pr_agent/ask_comment.png){width=512} +![Ask Comment](../assets/ask_comment.png){width=512} -![Ask](https://codium.ai/images/pr_agent/ask.png){width=512} +![Ask](../assets/ask.png){width=512} ## Ask lines @@ -21,7 +21,7 @@ You can run `/ask` on specific lines of code in the PR from the PR's diff view. - To select multiple lines, click on the '+' sign of the first line and then hold and drag to select the rest of the lines. - write `/ask "..."` in the comment box and press `Add single comment` button. -![Ask Line](https://codium.ai/images/pr_agent/Ask_line.png){width=512} +![Ask Line](../assets/Ask_line.png){width=512} Note that the tool does not have "memory" of previous questions, and answers each question independently. @@ -44,20 +44,20 @@ To get a direct link to an image, we recommend using the following scheme: 1\. First, post a comment that contains **only** the image: -![Ask image1](https://codium.ai/images/pr_agent/ask_images1.png){width=512} +![Ask image1](../assets/ask_images1.png){width=512} 2\. Quote reply to that comment: -![Ask image2](https://codium.ai/images/pr_agent/ask_images2.png){width=512} +![Ask image2](../assets/ask_images2.png){width=512} 3\. In the screen opened, type the question below the image: -![Ask image3](https://codium.ai/images/pr_agent/ask_images3.png){width=512} -![Ask image4](https://codium.ai/images/pr_agent/ask_images4.png){width=512} +![Ask image3](../assets/ask_images3.png){width=512} +![Ask image4](../assets/ask_images4.png){width=512} 4\. Post the comment, and receive the answer: -![Ask image5](https://codium.ai/images/pr_agent/ask_images5.png){width=512} +![Ask image5](../assets/ask_images5.png){width=512} See a full video tutorial [here](https://codium.ai/images/pr_agent/ask_image_video.mov) diff --git a/docs/docs/tools/describe.md b/docs/docs/tools/describe.md index 7cc1458ace..cca67d4e07 100644 --- a/docs/docs/tools/describe.md +++ b/docs/docs/tools/describe.md @@ -14,11 +14,11 @@ The tool can be triggered automatically every time a new PR is [opened](../usage Invoke the tool manually by commenting `/describe` on any PR: -![Describe comment](https://codium.ai/images/pr_agent/describe_comment.png){width=512} +![Describe comment](../assets/describe_comment.png){width=512} After ~30 seconds, the tool will generate a description for the PR: -![Describe New](https://codium.ai/images/pr_agent/describe_new.png){width=512} +![Describe New](../assets/describe_new.webp){width=512} If you want to edit [configurations](#configuration-options), add the relevant ones to the command: @@ -59,7 +59,7 @@ From our experience, there are two possible reasons: - When _updating_ PR descriptions, the `/describe` tool considers everything above the "PR Type" field as user content and will preserve it. Everything below this marker is treated as previously auto-generated content and will be replaced. -![Describe comment](https://codium.ai/images/pr_agent/pr_description_user_description.png){width=512} +![Describe comment](../assets/pr_description_user_description.png){width=512} ## Sequence Diagram Support The `/describe` tool includes a Mermaid sequence diagram showing component/function interactions. @@ -193,11 +193,11 @@ pr_agent:diagram The marker `pr_agent:type` will be replaced with the PR type, `pr_agent:summary` will be replaced with the PR summary, `pr_agent:walkthrough` will be replaced with the PR walkthrough, and `pr_agent:diagram` will be replaced with the sequence diagram (if enabled). -![Describe markers before](https://codium.ai/images/pr_agent/describe_markers_before.png){width=512} +![Describe markers before](../assets/describe_markers_before.png){width=512} becomes -![Describe markers after](https://codium.ai/images/pr_agent/describe_markers_after.png){width=512} +![Describe markers after](../assets/describe_markers_after.webp){width=512} **Configuration params**: @@ -258,7 +258,7 @@ Examples for custom labels: - ... The description should be comprehensive and detailed, indicating when to add the desired label. For example: -![Add native custom labels](https://codium.ai/images/pr_agent/add_native_custom_labels.png){width=768} +![Add native custom labels](../assets/add_native_custom_labels.webp){width=768} ## Usage Tips diff --git a/docs/docs/tools/generate_labels.md b/docs/docs/tools/generate_labels.md index ec6667a6f5..6ce1765238 100644 --- a/docs/docs/tools/generate_labels.md +++ b/docs/docs/tools/generate_labels.md @@ -10,13 +10,9 @@ It can be invoked manually by commenting on any PR: ## Example usage -Invoke the tool manually by commenting `/generate_labels` on any PR: +Invoke the tool manually by commenting `/generate_labels` on any PR. -![Generate Labels](https://codium.ai/images/pr_agent/generate_labels_comment.png){width=512} - -The tool will analyze the PR and add appropriate labels: - -![Generate Labels Result](https://codium.ai/images/pr_agent/generate_labels_result.png){width=512} +The tool will analyze the PR and add appropriate labels. ## Configuration options diff --git a/docs/docs/tools/help.md b/docs/docs/tools/help.md index 900587f772..54716269e9 100644 --- a/docs/docs/tools/help.md +++ b/docs/docs/tools/help.md @@ -13,9 +13,9 @@ It can be invoked manually by commenting on any PR: Invoke the `help` tool by commenting on a PR with: -![Help tool input](https://codium.ai/images/pr_agent/help1.png){width=750} +![Help tool input](../assets/help1.png){width=750} Response will include a list of available tools: -![Help tool output](https://codium.ai/images/pr_agent/help2.png){width=750} +![Help tool output](../assets/help2.webp){width=750} diff --git a/docs/docs/tools/help_docs.md b/docs/docs/tools/help_docs.md index 8a27c8c07c..d5a8ca7720 100644 --- a/docs/docs/tools/help_docs.md +++ b/docs/docs/tools/help_docs.md @@ -37,15 +37,15 @@ See more configuration options in the [Configuration options](#configuration-opt [//]: # (#### Asking a question about this repository:) -[//]: # (![help_docs on the documentation of this repository](https://codium.ai/images/pr_agent/help_docs_comment.png){width=512}) +[//]: # (![help_docs on the documentation of this repository](../assets/help_docs_comment.png){width=512}) **Asking a question about another repository** -![help_docs on the documentation of another repository](https://codium.ai/images/pr_agent/help_docs_comment_explicit_git.png){width=512} +![help_docs on the documentation of another repository](../assets/help_docs_comment_explicit_git.png){width=512} **Response**: -![help_docs response](https://codium.ai/images/pr_agent/help_docs_response.png){width=512} +![help_docs response](../assets/help_docs_response.png){width=512} ## Run automatically when a new issue is opened diff --git a/docs/docs/tools/improve.md b/docs/docs/tools/improve.md index 00396989ee..89a6621808 100644 --- a/docs/docs/tools/improve.md +++ b/docs/docs/tools/improve.md @@ -10,10 +10,10 @@ The tool can be triggered automatically every time a new PR is [opened](../usage ## How it looks === "Suggestions Overview" - ![code_suggestions_as_comment_closed](https://codium.ai/images/pr_agent/code_suggestions_as_comment_closed.png){width=512} + ![code_suggestions_as_comment_closed](../assets/code_suggestions_as_comment_closed.png){width=512} === "Selecting a specific suggestion" - ![code_suggestions_as_comment_open](https://codium.ai/images/pr_agent/code_suggestions_as_comment_open.png){width=512} + ![code_suggestions_as_comment_open](../assets/code_suggestions_as_comment_open.png){width=512} ___ @@ -45,7 +45,7 @@ When this fallback output succeeds, it is retained without an additional failure error propagation still follows `config.propagate_tool_errors`. If any individual retry succeeds, the existing partial-recovery behavior is preserved. -![improve](https://codium.ai/images/pr_agent/improve.png){width=512} +![improve](../assets/improve.png){width=512} ### Automatic triggering @@ -76,9 +76,9 @@ changes exits without calling the model. PR-Agent supports two modes for presenting code suggestions: -1) [Table](https://codium.ai/images/pr_agent/code_suggestions_as_comment_closed.png) mode +1) [Table](../assets/code_suggestions_as_comment_closed.png) mode -2) [Inline Committable](https://codium.ai/images/pr_agent/improve.png) code comments mode. +2) [Inline Committable](../assets/improve.png) code comments mode. The table format offers several key advantages: @@ -90,7 +90,7 @@ The table format offers several key advantages: Table mode is the default of PR-Agent, and is recommended approach for most users due to these benefits. -![code_suggestions_as_comment_closed.png](https://codium.ai/images/pr_agent/code_suggestions_as_comment_closed.png){width=512} +![code_suggestions_as_comment_closed.png](../assets/code_suggestions_as_comment_closed.png){width=512} Teams with specific preferences can enable committable code comments mode in their local configuration, or use [dual publishing mode](#dual-publishing-mode). @@ -122,25 +122,20 @@ Use triple quotes to write multi-line instructions. Use bullet points or numbers ### Best practices -`Platforms supported: GitHub, GitLab, Bitbucket` +`Repository context files supported: GitHub, GitLab, Gitea, Bitbucket, Azure DevOps` -!!! warning "Open-source PR-Agent" +The open-source PR-Agent package does not load `best_practices.md` automatically. To use a best-practices file, +add it to `config.repo_context_files`: - Automatic loading of `best_practices.md` is a Qodo Merge feature and is not available in the open-source - PR-Agent package. In the open-source package, add the file to `config.repo_context_files` instead: - - ```toml - [config] - repo_context_files = ["AGENTS.md", "best_practices.md"] - ``` - - This fallback supports GitHub, GitLab, Gitea, Bitbucket, and Azure DevOps. Repository context files are read - from the default branch by default and are limited by - `config.repo_context_max_lines` (500 lines by default). Set `config.repo_context_from_default_branch = false` - to read them from the pull request's target branch instead. Providers without repository file fetching log a - warning and skip this context. +```toml +[config] +repo_context_files = ["AGENTS.md", "best_practices.md"] +``` -Qodo Merge supports both simple and hierarchical best practices configurations to provide guidance to the AI model for generating relevant code suggestions. +This fallback supports GitHub, GitLab, Gitea, Bitbucket, and Azure DevOps. Repository context files are read +from the default branch by default and are limited by `config.repo_context_max_lines` (500 lines by default). +Set `config.repo_context_from_default_branch = false` to read them from the pull request's target branch instead. +Providers without repository file fetching log a warning and skip this context. ???- tip "Writing effective best practices files" @@ -149,10 +144,11 @@ Qodo Merge supports both simple and hierarchical best practices configurations t - Write clearly and concisely - Include brief code examples when helpful with before/after patterns - Focus on project-specific guidelines that will result in relevant suggestions you actually want to get - - Keep each file relatively short, under 800 lines, since: + - Keep each file short and focused, since: - AI models may not process effectively very long documents - Long files tend to contain generic guidelines already known to AI - - Maximum multiple file accumulated content is limited to 2000 lines. + - `config.repo_context_max_lines` limits the total rendered repository context, including wrapper and + truncation lines. - Use pattern-based structure rather than simple bullet points for better clarity ???- tip "Example of a best practices file" @@ -198,17 +194,11 @@ Qodo Merge supports both simple and hierarchical best practices configurations t return "" ``` -#### Local best practices in Qodo Merge - -For basic usage, create a `best_practices.md` file in your repository's root directory containing a list of best practices, coding standards, and guidelines specific to your repository. - -The AI model will use this `best_practices.md` file as a reference, and in case the PR code violates any of the guidelines, it will create additional suggestions, with a dedicated label: `Organization best practice`. - ### Combining 'extra instructions' and 'best practices' The `extra instructions` configuration is more related to the `improve` tool prompt. It can be used, for example, to avoid specific suggestions ("Don't suggest to add try-except block", "Ignore changes in toml files", ...) or to emphasize specific aspects or formats ("Answer in Japanese", "Give only short suggestions", ...) -In contrast, the `best_practices.md` file is a general guideline for the way code should be written in the repo. +In contrast, a `best_practices.md` file is a general guideline for the way code should be written in the repo. Using a combination of both can help the AI model to provide relevant and tailored suggestions. @@ -302,7 +292,7 @@ You can set the content of the checkbox text via: code_suggestions_self_review_text = "... (your text here) ..." ``` -![self_review_1](https://codium.ai/images/pr_agent/self_review_1.png){width=512} +![self_review_1](../assets/self_review_1.webp){width=512} !!! note "The checkbox is a visual marker only" @@ -417,6 +407,33 @@ for the authoritative default values. +## Recovering a partially failed analysis + +A partially failed analysis automatically tries the remaining configured +`config.fallback_models` for the chunks that raised an exception after other chunks succeeded, +with no extra configuration beyond `config.fallback_models` itself. +Successful predictions are retained in their original order. Recovery runs model by model, +after the preceding batch finishes, using the matching `openai.fallback_deployments` where configured, +and follows the effective fallback chain of the invocation, including a primary model selected by +[model routing](../usage-guide/changing_a_model.md#routing-small-pull-requests-to-a-cheaper-model). +A fallback model that recovers at least one chunk is marked as a fallback in the run +details, so the published run line still reports the primary model while the sticky fallback flag +stays set. + +Each remaining model is tried at most once per still-failed chunk at the tool level; the model +handler's existing retries and self-reflection calls can add provider requests. A chunk whose complete +generation prompt does not fit a fallback model's token budget is skipped for that model, without +truncating its context; the reserved headroom follows the output allowance the AI handler will +actually request for that model (`config.max_output_tokens` and related controls), falling back to +the fixed soft threshold when no specific limit is exposed. Later, larger models can still recover +it. Exhausted failures remain in the coverage notice. Cancellation propagates instead of starting +another retry. + +This is recovery within one invocation, not persisted resume across commands. Entirely failed initial +batches still use the existing whole-batch fallback. Valid empty suggestions are successful results; +responses rejected by the existing YAML parser retain its existing parse-failure behavior. +An ambiguous repeated model/deployment pair skips partial recovery rather than restarting the chain. + ## Understanding AI Code Suggestions - **AI Limitations:** AI models for code are getting better and better, but they are not flawless. Not all the suggestions will be perfect, and a user should not accept all of them automatically. Critical reading and judgment are required. Mistakes of the AI are rare but can happen, and it is usually quite easy for a human to spot them. diff --git a/docs/docs/tools/review.md b/docs/docs/tools/review.md index 5ff2b5ec4d..7ff51f00f4 100644 --- a/docs/docs/tools/review.md +++ b/docs/docs/tools/review.md @@ -18,11 +18,11 @@ Note that the main purpose of the `review` tool is to provide the **PR reviewer* Invoke the tool manually by commenting `/review` on any PR: -![review comment](https://codium.ai/images/pr_agent/review_comment.png){width=512} +![review comment](../assets/review_comment.png){width=512} After ~30 seconds, the tool will generate a review for the PR: -![review](https://codium.ai/images/pr_agent/review3.png){width=512} +![review](../assets/review3.png){width=512} If you want to edit [configurations](#configuration-options), add the relevant ones to the command: @@ -83,7 +83,7 @@ for the authoritative default values. persistent_finding_state - If set to true, PR-Agent persists structured review finding state across complete review runs, so findings can be resolved and reopened. Incremental and partial reviews do not resolve absent findings. Default is true. + If set to true, PR-Agent persists structured review finding state across complete review runs, so findings can be resolved and reopened. Resolved findings retain the original Markdown formatting of their description. Incremental and partial reviews do not resolve absent findings. Default is true. @@ -154,7 +154,11 @@ for the authoritative default values. require_ticket_analysis_review - If set to true, and the PR contains a GitHub or Jira ticket link, the tool will add a section that checks if the PR in fact fulfilled the ticket requirements. + + If set to true and ticket context is available, the tool adds a ticket-compliance section to the review comment + that checks whether the PR fulfills the ticket requirements. Supported sources include GitHub and GitLab issues, + linked Azure DevOps work items, Jira Cloud tickets, and Asana tasks. + require_risk_assessment @@ -213,11 +217,15 @@ for the authoritative default values. !!! tip "" - The `review` can tool automatically add labels to your Pull Requests: + The `review` tool can automatically add labels to your Pull Requests: - **`possible security issue`**: This label is applied if the tool detects a potential [security vulnerability](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/pr_reviewer_prompts.toml#L134) in the PR's code. This feedback is controlled by the 'enable_review_labels_security' flag (default is true). - **`review effort [x/5]`**: This label estimates the [effort](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/pr_reviewer_prompts.toml#L118) required to review the PR on a relative scale of 1 to 5, where 'x' represents the assessed effort. This feedback is controlled by the 'enable_review_labels_effort' flag (default is true). - - **`ticket compliance`**: Adds a label indicating code compliance level ("Fully compliant" | "PR Code Verified" | "Partially compliant" | "Not compliant") to any GitHub/Jira/Linea ticket linked in the PR. Controlled by the 'require_ticket_labels' flag (default: false). If 'require_no_ticket_labels' is also enabled, PRs without ticket links will receive a "No ticket found" label. + + Ticket compliance is reported in the review comment, not as a PR label. It is controlled by + `pr_reviewer.require_ticket_analysis_review` and requires available ticket context. The tool does not add + ticket-compliance labels or a "No ticket found" label. See + [Fetching ticket context](../core-abilities/fetching_ticket_context.md) for linking and authentication details. ### Auto-blocking PRs from being merged based on the generated labels @@ -226,7 +234,7 @@ for the authoritative default values. You can configure a CI/CD Action to prevent merging PRs with specific labels. For example, implement a dedicated [GitHub Action](https://medium.com/sequra-tech/quick-tip-block-pull-request-merge-using-labels-6cc326936221). - This approach helps ensure PRs with potential security issues or ticket compliance problems will not be merged without further review. + This approach helps ensure PRs with potential security issues will not be merged without further review. Since AI may make mistakes or lack complete context, use this feature judiciously. For flexibility, users with appropriate permissions can remove generated labels when necessary. When a label is removed, this action will be automatically documented in the PR discussion, clearly indicating it was a deliberate override by an authorized user to allow the merge. diff --git a/docs/docs/tools/similar_issues.md b/docs/docs/tools/similar_issues.md index 741c79938b..5d66719ed9 100644 --- a/docs/docs/tools/similar_issues.md +++ b/docs/docs/tools/similar_issues.md @@ -1,5 +1,7 @@ ## Overview +> **Note**: `/similar_issue` is an **experimental** feature. It works only on GitHub, carries a disproportionately large share of the project's dependency and configuration surface for a single-provider tool, and is therefore excluded from the v1 stability guarantees. Its backends are not equally exercised: the lancedb flow has no tests of its own, and no backend is tested against its real driver. + The similar issue tool retrieves the most similar issues to the current issue. It can be invoked manually by commenting on any PR: @@ -9,11 +11,11 @@ It can be invoked manually by commenting on any PR: ## Example usage -![similar_issue_original_issue](https://codium.ai/images/pr_agent/similar_issue_original_issue.png){width=768} +![similar_issue_original_issue](../assets/similar_issue_original_issue.png){width=768} -![similar_issue_comment](https://codium.ai/images/pr_agent/similar_issue_comment.png){width=768} +![similar_issue_comment](../assets/similar_issue_comment.png){width=768} -![similar_issue](https://codium.ai/images/pr_agent/similar_issue.png){width=768} +![similar_issue](../assets/similar_issue.png){width=768} Note that to perform retrieval, the `similar_issue` tool indexes all the repo previous issues (once). @@ -36,10 +38,31 @@ To use Pinecone with the `similar issue` tool, add these credentials to `.secret ``` [pinecone] api_key = "..." -environment = "..." +cloud = "aws" +region = "us-east-1" ``` -These parameters can be obtained by registering to [Pinecone](https://app.pinecone.io/?sessionType=signup/). +The `cloud` value must be one of `aws`, `gcp` or `azure`, and `region` must be an +availability region offered by that cloud. These parameters can be obtained by +registering to [Pinecone](https://app.pinecone.io/?sessionType=signup/). Note that the +tool uses Pinecone's serverless index API; the former `environment` setting from the +gcp-starter pod tier is no longer supported. + +`cloud` and `region` are only used when the index does not exist yet and needs to be +created. An existing index is opened by name and is never recreated, so moving an +existing deployment to the new configuration does not lose the stored vectors. + +!!! note "Backend coverage is uneven" + + No backend is exercised against its real driver: the `similar-issue` dependency group is + not installed in CI, so the pinecone tests run against a faked module and the qdrant tests + never construct a client. The lancedb flow has no tests of its own at all. + +!!! note "Default vector database" + + `vectordb` defaults to `lancedb`, which works with no external credentials. To use + qdrant or pinecone, set `vectordb = "qdrant"` or `vectordb = "pinecone"` under + `[pr_similar_issue]`. #### Qdrant Configuration @@ -60,7 +83,7 @@ vectordb = "qdrant" You can get a free managed Qdrant instance from [Qdrant Cloud](https://cloud.qdrant.io/). -Qdrant points are stored in a collection named `codium-ai-pr-agent-issues-v2`. +Qdrant points are stored in a collection named `codium-ai-pr-agent-issues-v2`, derived by appending a `-v2` suffix to the shared index name (`codium-ai-pr-agent-issues`). The suffix is an implementation detail of the Qdrant backend only; pinecone and lancedb use the unsuffixed name. !!! note "Upgrading an index created before the point-id fix" diff --git a/docs/docs/tools/update_changelog.md b/docs/docs/tools/update_changelog.md index 3064ba0855..53b8b0399d 100644 --- a/docs/docs/tools/update_changelog.md +++ b/docs/docs/tools/update_changelog.md @@ -9,9 +9,9 @@ It can be invoked manually by commenting on any PR: ## Example usage -![update_changelog_comment](https://codium.ai/images/pr_agent/update_changelog_comment.png){width=768} +![update_changelog_comment](../assets/update_changelog_comment.png){width=768} -![update_changelog](https://codium.ai/images/pr_agent/update_changelog.png){width=768} +![update_changelog](../assets/update_changelog.png){width=768} ## Configuration options diff --git a/docs/docs/usage-guide/EXAMPLE_BEST_PRACTICE.md b/docs/docs/usage-guide/EXAMPLE_BEST_PRACTICE.md deleted file mode 100644 index 993c0246db..0000000000 --- a/docs/docs/usage-guide/EXAMPLE_BEST_PRACTICE.md +++ /dev/null @@ -1,190 +0,0 @@ -## Recommend Python Best Practices - -This document outlines a series of recommended best practices for Python development. These guidelines aim to improve code quality, maintainability, and readability. - -### Imports - -Use `import` statements for packages and modules only, not for individual types, classes, or functions. - -#### Definition - -Reusability mechanism for sharing code from one module to another. - -#### Decision - -- Use `import x` for importing packages and modules. -- Use `from x import y` where `x` is the package prefix and `y` is the module name with no prefix. -- Use `from x import y as z` in any of the following circumstances: - - Two modules named `y` are to be imported. - - `y` conflicts with a top-level name defined in the current module. - - `y` conflicts with a common parameter name that is part of the public API (e.g., `features`). - - `y` is an inconveniently long name, or too generic in the context of your code -- Use `import y as z` only when `z` is a standard abbreviation (e.g., `import numpy as np`). - -For example the module `sound.effects.echo` may be imported as follows: - -``` -from sound.effects import echo -... -echo.EchoFilter(input, output, delay=0.7, atten=4) - -``` - -Do not use relative names in imports. Even if the module is in the same package, use the full package name. This helps prevent unintentionally importing a package twice. - -##### Exemptions - -Exemptions from this rule: - -- Symbols from the following modules are used to support static analysis and type checking: - - [`typing` module](https://google.github.io/styleguide/pyguide.html#typing-imports) - - [`collections.abc` module](https://google.github.io/styleguide/pyguide.html#typing-imports) - - [`typing_extensions` module](https://github.com/python/typing_extensions/blob/main/README.md) -- Redirects from the [six.moves module](https://six.readthedocs.io/#module-six.moves). - -### Packages - -Import each module using the full pathname location of the module. - -#### Decision - -All new code should import each module by its full package name. - -Imports should be as follows: - -``` -Yes: - # Reference absl.flags in code with the complete name (verbose). - import absl.flags - from doctor.who import jodie - - _FOO = absl.flags.DEFINE_string(...) - -``` - -``` -Yes: - # Reference flags in code with just the module name (common). - from absl import flags - from doctor.who import jodie - - _FOO = flags.DEFINE_string(...) - -``` - -_(assume this file lives in `doctor/who/` where `jodie.py` also exists)_ - -``` -No: - # Unclear what module the author wanted and what will be imported. The actual - # import behavior depends on external factors controlling sys.path. - # Which possible jodie module did the author intend to import? - import jodie - -``` - -The directory the main binary is located in should not be assumed to be in `sys.path` despite that happening in some environments. This being the case, code should assume that `import jodie` refers to a third-party or top-level package named `jodie`, not a local `jodie.py`. - -### Default Iterators and Operators - -Use default iterators and operators for types that support them, like lists, dictionaries, and files. - -#### Definition - -Container types, like dictionaries and lists, define default iterators and membership test operators (“in” and “not in”). - -#### Decision - -Use default iterators and operators for types that support them, like lists, dictionaries, and files. The built-in types define iterator methods, too. Prefer these methods to methods that return lists, except that you should not mutate a container while iterating over it. - -``` -Yes: for key in adict: ... - if obj in alist: ... - for line in afile: ... - for k, v in adict.items(): ... -``` - -``` -No: for key in adict.keys(): ... - for line in afile.readlines(): ... -``` - -### Lambda Functions - -Okay for one-liners. Prefer generator expressions over `map()` or `filter()` with a `lambda`. - -#### Decision - -Lambdas are allowed. If the code inside the lambda function spans multiple lines or is longer than 60-80 chars, it might be better to define it as a regular [nested function](https://google.github.io/styleguide/pyguide.html#lexical-scoping). - -For common operations like multiplication, use the functions from the `operator` module instead of lambda functions. For example, prefer `operator.mul` to `lambda x, y: x * y`. - -### Default Argument Values - -Okay in most cases. - -#### Definition - -You can specify values for variables at the end of a function’s parameter list, e.g., `def foo(a, b=0):`. If `foo` is called with only one argument, `b` is set to 0. If it is called with two arguments, `b` has the value of the second argument. - -#### Decision - -Okay to use with the following caveat: - -Do not use mutable objects as default values in the function or method definition. - -``` -Yes: def foo(a, b=None): - if b is None: - b = [] -Yes: def foo(a, b: Sequence | None = None): - if b is None: - b = [] -Yes: def foo(a, b: Sequence = ()): # Empty tuple OK since tuples are immutable. - ... -``` - -``` -from absl import flags -_FOO = flags.DEFINE_string(...) - -No: def foo(a, b=[]): - ... -No: def foo(a, b=time.time()): # Is `b` supposed to represent when this module was loaded? - ... -No: def foo(a, b=_FOO.value): # sys.argv has not yet been parsed... - ... -No: def foo(a, b: Mapping = {}): # Could still get passed to unchecked code. - ... -``` - -### True/False Evaluations - -Use the “implicit” false if possible, e.g., `if foo:` rather than `if foo != []:` - -### Lexical Scoping - -Okay to use. - -An example of the use of this feature is: - -``` -def get_adder(summand1: float) -> Callable[[float], float]: - """Returns a function that adds numbers to a given number.""" - def adder(summand2: float) -> float: - return summand1 + summand2 - - return adder -``` - -#### Decision - -Okay to use. - -### Threading - -Do not rely on the atomicity of built-in types. - -While Python’s built-in data types such as dictionaries appear to have atomic operations, there are corner cases where they aren’t atomic (e.g. if `__hash__` or `__eq__` are implemented as Python methods) and their atomicity should not be relied upon. Neither should you rely on atomic variable assignment (since this in turn depends on dictionaries). - -Use the `queue` module’s `Queue` data type as the preferred way to communicate data between threads. Otherwise, use the `threading` module and its locking primitives. Prefer condition variables and `threading.Condition` instead of using lower-level locks. diff --git a/docs/docs/usage-guide/additional_configurations.md b/docs/docs/usage-guide/additional_configurations.md index d0e6bb81f3..cae1b8f909 100644 --- a/docs/docs/usage-guide/additional_configurations.md +++ b/docs/docs/usage-guide/additional_configurations.md @@ -9,7 +9,7 @@ To print all the available configurations as a comment on your PR, you can use t /config ``` -![possible_config1](https://codium.ai/images/pr_agent/possible_config1.png){width=512} +![possible_config1](../assets/possible_config1.png){width=512} To view the **actual** configurations used for a specific tool, after all the user settings are applied, you can add for each tool a `--config.output_relevant_configurations=true` suffix. For example: @@ -20,7 +20,7 @@ For example: Will output an additional field showing the actual configurations used for the `improve` tool. -![possible_config2](https://codium.ai/images/pr_agent/possible_config2.png){width=512} +![possible_config2](../assets/possible_config2.png){width=512} ### Showing the agent run details @@ -336,7 +336,7 @@ Telemetry is disabled by default. To enable it, set in `configuration.toml`: ```toml [otel] is_enabled = true -exporter_type = "console" # "console", "otlp", or "none" +exporter_type = "console" # "console", "otlp", "prometheus", or "none" service_name = "pr-agent" environment = "development" # e.g. "development", "staging", "production" ``` @@ -362,6 +362,29 @@ otlp_protocol = "grpc" # default: "http" This is the recommended topology for fleets: point every PR-Agent instance at the same collector and aggregate there. Each process creates its own exporter connection; use the `service_name` and `environment` resource attributes to slice instances apart on the backend. +### Exposing native Prometheus metrics + +Instead of pushing to a collector, set `exporter_type = "prometheus"` to expose a native `GET /metrics` scrape endpoint on the gunicorn-served apps (`github_app`, `gitlab_webhook`, `azuredevops_server_webhook`, `gitea_app`). The command counter is translated into the Prometheus text format, and every gunicorn worker's values are merged at scrape time, so counters stay correct across the process workers: + +```toml +[otel] +exporter_type = "prometheus" +prometheus_multiproc_dir = "/tmp/pr-agent-prometheus" # shared, writable by every worker +``` + +1. The exporter is metrics-only: command spans are not exported in this mode. +2. `/metrics` is mounted only when this exporter is selected, so nothing is exposed by default. It does not depend on an OTLP collector or endpoint. +3. gunicorn registers and deregisters workers' state files automatically (`when_ready`/`child_exit`); a worker that dies mid-scrape leaves only a stale file, which is ignored once it is marked dead. +4. Metric families are created from the first data point's label set. Later attributes that do not fit the family are dropped, and missing ones are back-filled with an empty string, so a scrape never breaks on drifting label cardinality. +5. The exporter ships with PR-Agent (it depends on `prometheus-client`); no extra package is required. Scrape it like any exporter: + +```yaml +scrape_configs: + - job_name: pr-agent + static_configs: + - targets: ["pr-agent:3000"] +``` + Privacy controls (both off by default): - `include_pr_url = true` attaches PR URLs to spans. Off by default because URLs expose private repo names. @@ -373,6 +396,7 @@ Notes: - PR-Agent keeps its own OpenTelemetry providers and never registers the process-global one, so embedding PR-Agent in an application that already uses OpenTelemetry will not interfere with the host's telemetry. Pending spans and metrics are flushed automatically on process exit. - Each OTLP export call is bounded by `otlp_timeout` (default 3 seconds, retries included), so an unreachable collector cannot hang CLI exit or request completion. Raise it for slow collectors at the cost of longer worst-case stalls. - If `exporter_type = "otlp"` is set but no endpoint is configured, telemetry is disabled entirely (fail closed) — it never falls back to another exporter, so a missing secret cannot redirect telemetry into process logs. +- With `exporter_type = "prometheus"`, `prometheus_multiproc_dir` must be a shared directory writable by every worker; it defaults to `/tmp/pr-agent-prometheus`. In non-gunicorn (single-process) deployments the exporter works without it and simply serves the process's own registry. - **Serverless deployments** (e.g. the AWS Lambda webhooks) are supported: buffered spans and metrics are force-flushed at the end of every handled request, because frozen execution environments stop background export threads and are reaped without running exit handlers. No extra configuration is needed. ## Bringing per-repo context files to PR-Agent @@ -412,6 +436,41 @@ To bound how much of this context is sent to the model, `repo_context_max_lines` repo_context_max_lines = 500 ``` +### Context from sibling repositories + +The host operator must first approve repositories in the deployment configuration: + +```toml +[config] +repo_context_sibling_repos = ["my-group/library"] +repo_context_max_sibling_files = 5 +``` + +Approve only repositories whose contents may appear in reviews of consuming repositories. +This allowlist and the fetch limit cannot be changed by repository settings or comment arguments. +An empty allowlist disables sibling reads. + +A consuming repository can then select files in its `.pr_agent.toml` as structured +`repo_context_files` entries: + +```toml +[config] +repo_context_files = [ + "AGENTS.md", + {repo_id = "my-group/library", file_path = "src/api.py"}, +] +``` + +Use `owner/repository` on GitHub or a full project path on GitLab. GitLab also accepts numeric +project IDs as strings, provided the same identifier is in the host allowlist. Resolved repositories +must share the current repository's owning namespace: the GitHub owner or GitLab top-level group, +including projects in different subgroups. Renamed or redirected paths must be updated to their +canonical names in both settings. + +Sibling files always come from their default branch and share `repo_context_max_lines` with local +files. Private and internal repositories also require requester access. Comment arguments cannot +override `repo_context_files` at all; a repository's `.pr_agent.toml` may still set it. + ## Ignoring automatic commands in PRs PR-Agent allows you to automatically ignore certain PRs based on various criteria: @@ -517,19 +576,6 @@ ignore_language_framework = ['protobuf', ...] You can view the list of auto-generated file patterns in [`generated_code_ignore.toml`](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/generated_code_ignore.toml). Files matching these glob patterns will be automatically excluded from PR Agent analysis. -### Ignoring Tickets with Specific Labels - -When PR-Agent analyzes tickets (JIRA, GitHub Issues, GitLab Issues, etc.) referenced in your PR, you may want to exclude tickets that have certain labels from the analysis. This is useful for filtering out tickets marked as "ignore-compliance", "skip-review", or other labels that indicate the ticket should not be considered during PR review. - -To ignore tickets with specific labels, add the following to your `configuration.toml` file: - -```toml -[config] -ignore_ticket_labels = ["ignore-compliance", "skip-review", "wont-fix"] -``` - -Where `ignore_ticket_labels` is a list of label names that should be ignored during ticket analysis. - ### Restricted Mode When running PR-Agent with limited GitHub/GitLab permissions, set `restricted_mode` to `true` to gracefully skip operations that require elevated access (e.g., pushing changelog changes): diff --git a/docs/docs/usage-guide/automations_and_usage.md b/docs/docs/usage-guide/automations_and_usage.md index cadd1830e2..d674f857f5 100644 --- a/docs/docs/usage-guide/automations_and_usage.md +++ b/docs/docs/usage-guide/automations_and_usage.md @@ -34,6 +34,14 @@ This is useful for debugging or experimenting with different tools. 3. **git provider**: The [git_provider](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) field in a configuration file determines the GIT provider that will be used by PR-Agent. Currently, the following providers are supported: `github` **(default)**, `gitlab`, `bitbucket`, `azure`, `codecommit`, `local`, and `gitea`. +4. For scripts that need a failed request to return a non-zero process status, enable tool error propagation: + +```bash +python -m pr_agent.cli --pr_url= review --config.propagate_tool_errors=true +``` + +When a propagated tool error makes the request fail, the installed `pr-agent` command, `python -m pr_agent.cli`, and the customizable pip script exit with status 1. The default remains compatible with existing CLI behavior and exits with status 0; argparse parse and usage errors continue to exit with status 2. + ### CLI Health Check To verify that PR-Agent has been configured correctly, you can run this health check command from the repository root: @@ -93,8 +101,7 @@ When this parameter is set to `true`, PR-Agent will not run any automatic tools ### GitHub App !!! note "Configurations for PR-Agent" - PR-Agent for GitHub is an App, hosted by Codium. So all the instructions below are relevant for PR-Agent users. - Same goes for [GitLab webhook](#gitlab-webhook) and [BitBucket App](#bitbucket-app) sections. + These settings apply to self-hosted GitHub App, GitLab webhook, and Bitbucket App deployments. #### GitHub app automatic tools when a new PR is opened diff --git a/docs/docs/usage-guide/changing_a_model.md b/docs/docs/usage-guide/changing_a_model.md index 15b1d7a7fc..20fd9b97c5 100644 --- a/docs/docs/usage-guide/changing_a_model.md +++ b/docs/docs/usage-guide/changing_a_model.md @@ -20,6 +20,11 @@ You can give parameters via a configuration file, or from environment variables. See [litellm documentation](https://litellm.vercel.app/docs/proxy/quick_start#supported-llms) for the environment variables needed per model, as they may vary and change over time. Our documentation per-model may not always be up-to-date with the latest changes. Failing to set the needed keys of a specific model will usually result in litellm not identifying the model type, and failing to utilize it. +!!! warning "Credential isolation boundary" + PR-Agent captures request settings and supported provider environment values per handler. Deployment-owned LiteLLM secret managers are outside this request-isolation boundary; their credentials are not snapshotted by PR-Agent. Keep process environment variables and LiteLLM globals stable while requests run. The handler does not isolate arbitrary changes made by embedding applications during a request. Use separate processes when workloads require different deployment-owned credential sources or mutable global authentication and routing state. + + Process-wide routing fallbacks such as `litellm.api_base`, `litellm.api_version`, `litellm.organization`, `litellm.vertex_project`, and `litellm.vertex_location` can be rejected even when unchanged. Use the corresponding PR-Agent settings (`OPENAI.API_BASE`, `OPENAI.API_VERSION`, `OPENAI.ORG`, `VERTEXAI.VERTEX_PROJECT`, and `VERTEXAI.VERTEX_LOCATION`) or supported provider environment variables, and clear the corresponding LiteLLM globals even when they match the intended routing. Use `LITELLM.EXTRA_HEADERS` instead of `litellm.headers`. + ### OpenAI like API To use an OpenAI like API, set the following in your `.secrets.toml` file: @@ -43,7 +48,7 @@ To reduce costs for non-urgent/background tasks, enable Flex Processing: ```toml [litellm] -extra_body='{"processing_mode": "flex"}' +extra_body='{"service_tier": "flex"}' ``` See [OpenAI Flex Processing docs](https://platform.openai.com/docs/guides/flex-processing) for details. @@ -79,6 +84,10 @@ tenant_id = "" # Your Azure AD tenant ID api_base = "" # Your Azure OpenAI service base URL (e.g., https://openai.xyz.com/) ``` +The request-local Azure OIDC bridge captures `AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_AUTHORITY_HOST`, `AZURE_SCOPE`, and competing `AZURE_CLIENT_SECRET` / `AZURE_USERNAME` / `AZURE_PASSWORD` settings when the handler is initialized. It preserves native authentication precedence while binding companion credentials and their cache identity to the captured authority; an absent authority uses the Azure public-cloud default, while an empty authority is rejected for companion credentials. LiteLLM's native assertion and secret-selector resolvers still control updates and source selection; their underlying deployment-owned credential chains are not request-isolated. Keep those sources scoped to the intended workload identity, or use separate processes for distinct identities. + +For native Azure SDK routes other than Cloudflare gateways, ordinary AD tokens use the same captured companion settings and SDK client-cache isolation. Azure Responses routes also bind ordinary AD companion selection to the captured settings. Complete companion credentials can also be selected without an initial AD token, including on native Azure AI raw HTTP routes. Captured companion providers retain callable token refresh and native authentication precedence. If SDK initialization or Responses token resolution reaches LiteLLM's implicit credential discovery with `enable_azure_ad_token_refresh=True`, the request is rejected: that fallback re-reads process-wide identity settings. Configure complete client-secret or username/password companion credentials instead. Managed identity, certificate, and default credential-chain discovery through this fallback are intentionally unsupported; an already cached, isolated SDK client can still be reused without entering discovery. + Passing custom headers to the underlying LLM Model API can be done by setting extra_headers parameter to litellm. ```toml @@ -88,6 +97,8 @@ extra_headers='{"projectId": "", ...}') #The value of thi This enables users to pass authorization tokens or API keys, when routing requests through an API management gateway. +Requests that would otherwise inherit non-empty process-wide `litellm.headers` are rejected, even for non-authentication headers; configure headers through `LITELLM.EXTRA_HEADERS` instead. + ### Ollama You can run models locally through either [VLLM](https://docs.litellm.ai/docs/providers/vllm) or [Ollama](https://docs.litellm.ai/docs/providers/ollama) @@ -105,7 +116,7 @@ duplicate_examples=true # will duplicate the examples in the prompt, to help the api_base = "http://localhost:11434" # or whatever port you're running Ollama on ``` -By default, Ollama uses a context window size of 2048 tokens. In most cases this is not enough to cover pr-agent prompt and pull-request diff. Context window size can be overridden with the `OLLAMA_CONTEXT_LENGTH` environment variable. For example, to set the default context length to 8K, use: `OLLAMA_CONTEXT_LENGTH=8192 ollama serve`. More information you can find on the [official ollama faq](https://github.com/ollama/ollama/blob/main/docs/faq.md#how-can-i-specify-the-context-window-size). +By default, Ollama uses a context window size of 2048 tokens. In most cases this is not enough to cover pr-agent prompt and pull-request diff. Context window size can be overridden with the `OLLAMA_CONTEXT_LENGTH` environment variable. For example, to set the default context length to 8K, use: `OLLAMA_CONTEXT_LENGTH=8192 ollama serve`. More information you can find on the [official ollama faq](https://docs.ollama.com/faq#how-can-i-specify-the-context-window-size). Please note that the `custom_model_max_tokens` setting should be configured in accordance with the `OLLAMA_CONTEXT_LENGTH`. Failure to do so may result in unexpected model output. @@ -218,6 +229,12 @@ Your [application default credentials](https://cloud.google.com/docs/authenticat If you do want to set explicit credentials, then you can use the `GOOGLE_APPLICATION_CREDENTIALS` environment variable set to a path to a json credentials file. +Each handler captures the selected Cloud SDK ADC file and its resource project configuration, including `CLOUDSDK_CONFIG`, the active named configuration, and `CLOUDSDK_CORE_PROJECT`. Later changes apply to new handlers, not existing ones. The quota project remains separate from the resource project. When no ADC file exists at initialization, the handler retains managed-runtime discovery rather than adopting a file created later. + +AWS-backed Vertex workload identity federation captures environment credentials and region per handler. Different captured identities use separate credential caches even when they share the same WIF configuration; identical snapshots can reuse a cache entry. Metadata-backed credentials continue refreshing through the captured metadata source. + +Vertex external-account credentials with an executable source are intentionally unsupported, including when `GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES=1`. The helper and its cached output can change identity after a handler captures its configuration, so Vertex requests reject this source before using it for authentication. Use a non-executable credential source configured for the intended identity instead; normal credential refresh remains enabled for supported sources. + ### Google AI Studio To use [Google AI Studio](https://aistudio.google.com/) models, set the relevant models in the configuration section of the configuration file: @@ -299,6 +316,41 @@ Set `AWS_USE_IMDS=true` in the environment. PR-Agent will resolve credentials vi | EKS pod with IRSA | Web identity token + STS | | Lambda function | Runtime-injected credentials | +Credential discovery runs synchronously when the handler is initialized. Before each SigV4 call, PR-Agent +refreshes credentials synchronously through the same boto3 credentials object and passes a request-local snapshot +to LiteLLM, without writing credentials into the process environment. AWS calls using this provider chain are +serialized within a handler, including any static-credential retry. Discovery and refresh can block the event loop. + +The same opt-in is required for other boto3 provider-chain sources, including `AWS_PROFILE` and shared credentials +files. LiteLLM-specific `AWS_PROFILE_NAME` and `AWS_ROLE_NAME` selectors are not supported because they can override +request-local credentials; unset them and use `AWS_USE_IMDS=true` instead. + +When upgrading from implicit LiteLLM credential-chain discovery, set `AWS_USE_IMDS=true` explicitly. Without this +opt-in, PR-Agent only uses complete static credentials from settings or environment credentials captured when the +handler is initialized; it does not ask LiteLLM to discover a role or profile at request time. + +For classic Bedrock, a supported model ARN supplies the request region ahead of environment/settings regions, including during static-credential retries. Converse also uses a captured `litellm.model_id` ARN, or a region/model path when `model_id` is absent; Invoke does not derive its region from the separate `model_id`. Otherwise, without `AWS_USE_IMDS=true`, set `aws.AWS_REGION_NAME`, `AWS_REGION_NAME`, `AWS_REGION`, or `AWS_DEFAULT_REGION`. Complete environment credentials alone do not enable region discovery from boto3 profiles or LiteLLM's default region. Static credentials in `[aws]` still require `AWS_REGION_NAME`. + +For Bedrock Mantle without IMDS or complete static credentials, the region comes from `BEDROCK_MANTLE_REGION`, `AWS_REGION_NAME`, `aws.AWS_REGION_NAME`, or `AWS_REGION`, then defaults to `us-east-1`; `AWS_DEFAULT_REGION` alone does not change that default. A `BEDROCK_MANTLE_API_BASE` on the standard host `https://bedrock-mantle..api.aws` supplies the region and overrides these sources; a custom endpoint host does not. Complete static credentials and opted-in boto3 discovery retain the AWS region policy described above. + +Without `AWS_USE_IMDS=true`, environment authentication follows LiteLLM's `AWS_SESSION_TOKEN` selection. The legacy `AWS_SECURITY_TOKEN` alias is recognized only by the opted-in boto3 chain, which prefers it over `AWS_SESSION_TOKEN` when both are nonempty. + +Bedrock bearer tokens must come from handler-captured credentials rather than process-wide LiteLLM fallback, and `AWS_BEARER_TOKEN_BEDROCK` must be unset for `sagemaker_chat` and `sagemaker_nova` routes. + +When resolving credentials through profiles with `AWS_USE_IMDS=true`, PR-Agent does not execute +`credential_process` in the selected profile or its `source_profile` chain. If such a process is configured, +PR-Agent uses complete static credentials from `[aws]` (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and +`AWS_REGION_NAME`, plus `AWS_SESSION_TOKEN` when required) instead; without them, credential resolution fails. +This restriction does not apply when using complete environment credentials captured at handler initialization. + +Workload token files selected by `AWS_WEB_IDENTITY_TOKEN_FILE`, profile `web_identity_token_file`, or +`AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE` must be controlled by the deployment, not switched between tenants. +PR-Agent retains native token reload and credential refresh from the selected source; it does not freeze token +contents or detect a different identity replacing the contents at the same path. Use separately controlled +credential sources and appropriate process/container isolation for distinct workload identities, or supply +explicit request-local credentials. Request-local configuration is not an isolation boundary against arbitrary +code running in the same process. + Minimal GitHub Actions workflow (no AWS secret keys required): ```yaml @@ -321,7 +373,7 @@ The IAM role must have `bedrock:InvokeModel` permission on the target model ARN, } ``` -If you also configure static keys in `[aws]`, they serve as an automatic fallback: if the ambient credentials fail a Bedrock call (e.g., the role lacks `bedrock:InvokeModel`), PR-Agent retries with the static keys and logs a warning. +If you also configure static keys in `[aws]`, they serve as an automatic fallback when ambient credentials cannot be resolved or a SigV4 call using the opted-in provider chain raises an API error other than a rate-limit error. This applies to Bedrock, Bedrock Mantle, and SageMaker routes and preserves the existing fallback behavior for IAM authorization and connection failures, using only the static credentials captured by the handler. #### Custom Inference Profiles @@ -343,6 +395,39 @@ model_id = "your-application-inference-profile-arn" The `litellm.model_id` parameter applies only to classic `bedrock/` calls made through the `bedrock-runtime` APIs. It does not apply to `bedrock_mantle/`; for cost allocation with the Mantle Chat Completions and Responses APIs, use [Amazon Bedrock Projects](https://docs.aws.amazon.com/bedrock/latest/userguide/cost-mgmt-projects.html). +#### Claude 5 thinking with an application inference profile ARN + +Claude Sonnet 5 on Bedrock is invoked through an inference profile rather than a direct +foundation-model id. When that profile is an application inference profile, its ARN is an +opaque value that carries no model name. Thinking configuration in PR-Agent is gated on +recognizing the model, so the opaque ARN can never match: `enable_claude_adaptive_thinking` +requires a recognized Claude 5 model name in the id, and `enable_claude_extended_thinking` +requires exact membership in `claude_extended_thinking_models`. PR-Agent logs a warning in +that case, so the unconfigured state is no longer silent. + +Address the model by name and pass the profile ARN through `litellm.model_id`, which is what +the invocation actually uses: + +```toml +[config] # in configuration.toml +model = "bedrock/converse/eu.anthropic.claude-sonnet-5" +fallback_models = ["bedrock/converse/eu.anthropic.claude-sonnet-5"] +enable_claude_adaptive_thinking = true # requires a recognizable claude 5 model name in `model` + +[litellm] +model_id = "arn:aws:bedrock:eu-central-1::application-inference-profile/" +``` + +Cost attribution is preserved through the application inference profile, and because `model` +is the named id, the adaptive-thinking payload is applied and kept intact. + +Two caveats. First, ARNs only fail the detection when the suffix is opaque: an ARN that +embeds the model family, for example `...:inference-profile/us.anthropic.claude-sonnet-5`, +normalises to a string the adaptive regex does match. The miss is specific to application +inference profiles with an opaque hex suffix. Second, `litellm.model_id` is a single global +value applied to every model whose id contains `bedrock/`, so this configuration cannot point +different models at different profiles within one fallback chain without per-call handling. + #### Using a Custom VPC Endpoint (PrivateLink) To route Bedrock traffic through a VPC interface endpoint instead of the public `bedrock-runtime` endpoint, set `AWS_BEDROCK_RUNTIME_ENDPOINT` either as an environment variable or in `[aws]`: @@ -508,6 +593,8 @@ api_base = "https://adb-xxxx.azuredatabricks.net/serving-endpoints" # your works The model name after the `databricks/` prefix is the name of your serving endpoint. See LiteLLM's [Databricks provider docs](https://docs.litellm.ai/docs/providers/databricks) for details. +The configured PAT and endpoint are captured per handler. When both `DATABRICKS_CLIENT_ID` and `DATABRICKS_CLIENT_SECRET` are set, LiteLLM performs an OAuth M2M exchange even if a PAT is supplied. An exchange failure can abort the request; after a successful exchange, the supplied PAT replaces the OAuth token in the final Authorization header. Without OAuth M2M credentials or a PAT, optional Databricks SDK authentication remains available. These deployment-owned credentials and profile selection are not request-isolated. Do not change those sources between tenants in a shared process; use request-specific PATs with OAuth M2M credentials unset, or separate processes for distinct workload identities. + ### Openrouter To use model from Openrouter, for example, set: @@ -554,7 +641,7 @@ For `openrouter/...` models you can optionally restrict which upstream providers # max_tokens = 16000 # hard cap on completion tokens for the request ``` -`provider_only` and `reasoning_effort = "none"` are useful to pin a specific provider and to bound the cost of reasoning models. Because Openrouter treats effort and token budgets as mutually exclusive, an explicit Openrouter-specific `"none"` keeps reasoning disabled when the model supports disabling it. Grok 4.5/4.6 clamp `"none"` before precedence is applied, so a positive budget wins there; otherwise a positive `reasoning_max_tokens` value takes precedence over the global effort and other Openrouter-specific values. Invalid Openrouter-specific effort values are warned about and treated as unset, so registered reasoning models fall back to `config.reasoning_effort`. Openrouter normalizes `"max"` to `"xhigh"` in this path to match LiteLLM 1.98.0. Supported effort values vary by model, and models whose metadata marks reasoning as mandatory reject `"none"`. For Anthropic models using a reasoning budget, set the effective output `max_tokens` higher than `reasoning_max_tokens` so the final answer has output headroom. See the Openrouter [provider routing](https://openrouter.ai/docs/guides/routing/provider-selection) and [reasoning tokens](https://openrouter.ai/docs/guides/best-practices/reasoning-tokens) docs. +`provider_only` and `reasoning_effort = "none"` are useful to pin a specific provider and to bound the cost of reasoning models. Because Openrouter treats effort and token budgets as mutually exclusive, an explicit Openrouter-specific `"none"` keeps reasoning disabled when the model supports disabling it. Grok 4.5/4.6 clamp `"none"` before precedence is applied, so a positive budget wins there; otherwise a positive `reasoning_max_tokens` value takes precedence over the global effort and other Openrouter-specific values. Invalid Openrouter-specific effort values are warned about and treated as unset, so registered reasoning models fall back to `config.reasoning_effort`. Openrouter normalizes `"max"` to `"xhigh"` in this path for LiteLLM/OpenRouter compatibility. Supported effort values vary by model, and models whose metadata marks reasoning as mandatory reject `"none"`. For Anthropic models using a reasoning budget, set the effective output `max_tokens` higher than `reasoning_max_tokens` so the final answer has output headroom. See the Openrouter [provider routing](https://openrouter.ai/docs/guides/routing/provider-selection) and [reasoning tokens](https://openrouter.ai/docs/guides/best-practices/reasoning-tokens) docs. ### OrcaRouter @@ -615,6 +702,57 @@ Create the credential per branch in the [Neon Console](https://console.neon.tech !!! note "Chat completions only" Some model IDs in Neon's catalog are served through the OpenAI Responses API, which Neon exposes under `/openai/v1` instead of `/v1`. The configuration above points at the chat-completions endpoint, so it cannot reach those models. Neon also documents a few models that return `message.content` as an array of typed blocks rather than a string, and PR-Agent reads the reply as a string. +### GitHub Copilot + +Models under an active GitHub Copilot subscription are available through litellm's `github_copilot` provider, which authenticates as your GitHub identity rather than with a dedicated API key: + +```toml +[config] +model = "github_copilot/gpt-4o" +fallback_models = ["github_copilot/gpt-4.1"] +``` + +The GitHub identity behind the model needs an active Copilot subscription. The token budget for a Copilot model is resolved automatically from litellm's model metadata (verified against the pinned litellm 1.100.0), so `custom_model_max_tokens` is not required. However, `get_max_tokens` clamps the effective window to `config.max_model_tokens`, which defaults to 32000. To use the full context window of the model (e.g., 64000 for gpt-4o, 128000 for gpt-4.1), raise `config.max_model_tokens` accordingly. + +Authentication uses the [GitHub Copilot provider](https://docs.litellm.ai/docs/providers/github_copilot) flow: + +1. litellm first looks for a pre-seeded GitHub access token in `access-token` under `GITHUB_COPILOT_TOKEN_DIR` (default `~/.config/litellm/github_copilot`; the filename is overridable with `GITHUB_COPILOT_ACCESS_TOKEN_FILE`). In a CI runner, write that token before the job runs - for example, mount a secret into the directory or point the variable at a mounted secret directory - and the flow never becomes interactive. +2. Only when the file is missing or empty does litellm fall back to the interactive device-code flow (`POST https://github.com/login/device/code`, up to three attempts), which does not suit unattended runners. +3. The Copilot API key (`api-key.json` in the same directory) is refreshed automatically against `https://api.github.com/copilot_internal/v2/token`, using the pre-seeded access token. + +Whether Copilot's terms permit this programmatic use is a question for GitHub rather than a guarantee this project can make, so confirm before relying on the route. + +### Atlas Cloud + +[Atlas Cloud](https://www.atlascloud.ai) is an OpenAI-compatible inference platform. It needs no provider-specific code in PR-Agent: the `openai/` prefix routes the request to Atlas's base URL through litellm's OpenAI-compatible path, the same way [OrcaRouter](#orcarouter) and the [Neon AI Gateway](#neon-ai-gateway) are handled. + +To use a model served by Atlas Cloud, set: + +```toml +[config] # in configuration.toml +model = "openai/deepseek-ai/deepseek-v4-pro" +fallback_models = ["openai/deepseek-ai/deepseek-v4-pro"] +custom_model_max_tokens = 1048000 # the context window Atlas publishes for the model + +[openai] # in .secrets.toml +api_base = "https://api.atlascloud.ai/v1" +key = "..." # your Atlas Cloud api key +``` + +or use the environment variables (make sure to use double underscores `__`): + +```bash +OPENAI__API_BASE=https://api.atlascloud.ai/v1 +OPENAI__KEY=... +``` + +(you can obtain an Atlas Cloud API key from the [console](https://www.atlascloud.ai/console)) + +Keep the `openai/` prefix on the model name, whichever Atlas model ID you use (`openai/deepseek-ai/deepseek-v4-pro`, `openai/zai-org/glm-5`, `openai/moonshotai/kimi-k2.6`, ...): the prefix routes the request through litellm's OpenAI-compatible path. A prefixed name is not in the `MAX_TOKENS` table [here](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/algo/__init__.py), so you also have to set `custom_model_max_tokens`. Take the value from Atlas's [model catalog](https://www.atlascloud.ai/models). + +!!! note "Reasoning models need output headroom" + Several Atlas models are reasoning models that spend completion tokens on a hidden chain of thought before writing the answer. `deepseek-ai/deepseek-v4-pro` with `max_tokens = 16` returns `finish_reason = "length"` and an **empty** `message.content` — all 16 completion tokens were reasoning tokens. If a tool comes back blank, raise the output budget rather than assuming the request failed. Non-reasoning IDs such as `deepseek-ai/DeepSeek-V3.1` are unaffected. + ### Custom models If the relevant model doesn't appear [here](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/algo/__init__.py), you can still use it as a custom model: @@ -693,6 +831,11 @@ built-in defaults. for it and logs a warning rather than sending a request the provider would reject — use `enable_claude_adaptive_thinking` for those models instead. +Both thinking gates only fire when the model id itself is recognizable: an opaque id such as a +Bedrock application inference profile ARN matches neither gate, and PR-Agent logs a warning +instead of silently sending nothing. See [Claude 5 thinking with an application inference +profile ARN](#claude-5-thinking-with-an-application-inference-profile-arn). + ## Output token limit ```toml diff --git a/docs/docs/usage-guide/configuration_options.md b/docs/docs/usage-guide/configuration_options.md index 7d0521b785..1f356c1a96 100644 --- a/docs/docs/usage-guide/configuration_options.md +++ b/docs/docs/usage-guide/configuration_options.md @@ -8,7 +8,7 @@ There are three main ways to set persistent configurations: In terms of precedence, local configurations will override global configurations, and global configurations will override an external configuration URL. -For a list of all possible configurations, see the [configuration options](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) page. +For a list of all possible configurations, see the [configuration options](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) page, or the rendered [Configuration Reference](./configuration_reference.md) which lists every option grouped by section. In addition to general configuration options, each tool has its own configurations. For example, the `review` tool will use parameters from the [pr_reviewer](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) section in the configuration file. !!! tip "Tip1: Edit only what you need" @@ -89,11 +89,11 @@ Loading the global settings file is controlled by the `use_global_settings_file` use_global_settings_file = false ``` -For example, in the GitHub organization `qodo-ai`: +For example, in a GitHub organization named `my-org`: -- The file [`https://github.com/the-pr-agent/pr-agent-settings/.pr_agent.toml`](https://github.com/the-pr-agent/pr-agent-settings/blob/main/.pr_agent.toml) serves as a global configuration file for all the repos in the GitHub organization `qodo-ai`. +- The file `my-org/pr-agent-settings/.pr_agent.toml` (read from that repository's default branch) serves as a global configuration file for all the repos in the organization. -- The repo [`https://github.com/the-pr-agent/pr-agent`](https://github.com/the-pr-agent/pr-agent/blob/main/.pr_agent.toml) inherits the global configuration file from `pr-agent-settings`. +- A repository such as `my-org/my-repo` inherits that global configuration file, and may override any of its values in its own `.pr_agent.toml`. ## Project/Group level configuration file @@ -117,7 +117,7 @@ Create a dedicated project to hold a global configuration file that affects all 1. Create a new project with both the name and key: PR_AGENT_SETTINGS. 2. Inside the PR_AGENT_SETTINGS project, create a repository named pr-agent-settings. 3. In this repository, add a `.pr_agent.toml` configuration file—structured similarly to the global configuration file described above. -4. Optionally, you can add organizational-level [global best practices](../tools/improve.md#global-hierarchical-best-practices). +4. Optionally, you can add repository context files with shared instructions or best practices. Repositories across your entire Bitbucket organization will inherit the configuration from this file. diff --git a/docs/docs/usage-guide/configuration_reference.md b/docs/docs/usage-guide/configuration_reference.md new file mode 100644 index 0000000000..17127665f4 --- /dev/null +++ b/docs/docs/usage-guide/configuration_reference.md @@ -0,0 +1,543 @@ +# Configuration Reference + +> This page is **auto-generated** and should not be edited by hand. +> Regenerate it from the [TOML source](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) with: +> +> ```bash +> python scripts/generate_config_reference.py +> ``` + +Every configuration option PR-Agent supports, grouped by section. The [configuration.toml](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) +file is the single source of truth for defaults and inline comments; this page renders the same +list for easy searching and linking. + +Rows with an empty **Description** are keys whose TOML entry carries no explanatory comment yet. +They are listed deliberately rather than hidden, so the gaps double as the documentation +to-do list. + +## `[config]` + +**models** + +| Key | Default | Description | +| --- | --- | --- | +| `model` | "gpt-5.6" | | +| `fallback_models` | ["gpt-5.6-terra"] | | +**CLI** + +| Key | Default | Description | +| --- | --- | --- | +| `git_provider` | "github" | | +| `publish_output` | true | | +| `publish_output_progress` | true | | +| `progress_gif_url` | "" | optional, override for the progress loading gif url (example: 'https://.../spinner.gif'). | +| `progress_gif_width` | 48 | optional, width (in px) of the progress loading gif. | +| `verbosity_level` | 0 | 0,1,2 | +| `use_extra_bad_extensions` | false | | +**Log** + +| Key | Default | Description | +| --- | --- | --- | +| `log_level` | "DEBUG" | | +**Configurations** + +| Key | Default | Description | +| --- | --- | --- | +| `use_repo_settings_file` | true | | +| `use_global_settings_file` | true | | +| `enable_per_directory_settings` | false | when true, merge per-directory .pr_agent.toml files found by walking up from the PR's changed files (monorepo support). Adds bounded recursive tree discovery per MR; per-directory files may only override non-critical sections (see REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS). Nearest (deepest) directory wins on shared keys; equal-depth siblings resolve to the lexicographically-last path; when more files match than the per_directory_settings_max_files cap, shallower files are applied first and a partially capped depth keeps its later-path (winning) siblings; any overlap is logged as a warning. | +| `per_directory_settings_max_files` | 20 | hard ceiling on the number of per-directory .pr_agent.toml files applied per MR (deeper configs beyond the cap are skipped with a warning) | +| `per_directory_settings_max_tree_pages` | 10 | maximum GitLab recursive-tree pages (100 entries each); skip nested settings if discovery is incomplete. Root/host-controlled, independent of the settings-file cap. | +| `extra_config_url` | "" | optional URL or path to an additional .pr_agent.toml merged before the repo-local config; also settable via --extra_config_url or PR_AGENT_EXTRA_CONFIG_URL. See docs/docs/usage-guide/configuration_options.md#external-configuration-url. | +| `disable_auto_feedback` | false | | +| `enable_auto_approval` | false | when true, /review may auto-approve a PR via auto_approve_logic(); that caller is currently commented out | +| `ai_timeout` | 120 | 2 minutes | +| `retry_same_model_on_timeout` | true | when false, a timed-out call is not retried on the same model and moves on to fallback_models | +| `skip_keys` | [] | | +| `custom_reasoning_model` | false | when true, disables system messages and temperature controls for models that don't support chat-style inputs | +| `response_language` | "en-US" | Language locales code for PR responses in ISO 3166 and ISO 639 format (e.g., "en-US", "it-IT", "zh-CN", ...) | +| `repo_context_files` | ["AGENTS.md"] | Repository-relative files (e.g. AGENTS.md, CLAUDE.md) to include as AI prompt context; set to [] to disable local context. A structured entry {"repo_id" = ..., "file_path" = ...} selects a sibling default-branch file from the same namespace/owner; repo_id must be in the host-issued repo_context_sibling_repos allowlist below. Reads use the sibling default branch and share repo_context_max_lines; repository settings may select entries, comment arguments cannot override this key | +| `repo_context_from_default_branch` | true | Read repo context files from the repository default branch (trusts only default-branch content). Set to false to read from the PR target branch instead. | +| `repo_context_max_lines` | 500 | Maximum total rendered lines for repo context, including wrapper tags | +| `repo_context_sibling_repos` | [] | Host-only list of approved sibling repository identifiers (GitHub owner/repo, GitLab group/project or numeric ID strings) that repo_context_files sibling entries may select. Empty disables sibling reads. Approve only repositories whose content may be disclosed in consuming PRs, because the actor check bounds who triggers a read, not who chose the target or where the output lands. Repository settings and comment arguments cannot change this list. Canonical identities and owning namespaces are checked after resolution | +| `repo_context_max_sibling_files` | 5 | Maximum number of sibling-repository files fetched per repo-context build. The fetch count is bounded separately from repo_context_max_lines so selected sibling files cannot trigger an unbounded number of cross-repository calls; sibling files still compete for the repo_context_max_lines budget. Host-only (cannot be raised by a repository's .pr_agent.toml or a comment command) and clamped to a hard ceiling of 20 fetches per build. | +**token limits** + +| Key | Default | Description | +| --- | --- | --- | +| `max_description_tokens` | 500 | | +| `max_commits_tokens` | 500 | | +| `max_model_tokens` | 32000 | Limits the maximum number of tokens that can be used by any model, regardless of the model's default capabilities. | +| `custom_model_max_tokens` | -1 | for models not in the default list | +| `max_output_tokens` | 0 | 0 = unset (the provider's own default applies) | +| `model_token_count_estimate_factor` | 0.3 | factor to increase the token count estimate, in order to reduce likelihood of model failure due to too many tokens - applicable only when requesting an accurate estimate. | +**patch extension logic** + +| Key | Default | Description | +| --- | --- | --- | +| `patch_extension_skip_types` | [".md", ".txt"] | | +| `allow_dynamic_context` | true | | +| `max_extra_lines_before_dynamic_context` | 10 | will try to include up to 10 extra lines before the hunk in the patch, until we reach an enclosing function or class | +| `patch_extra_lines_before` | 5 | Number of extra lines (+3 default ones) to include before each hunk in the patch | +| `patch_extra_lines_after` | 1 | Number of extra lines (+3 default ones) to include after each hunk in the patch | +| `secret_provider` | "" | "" (disabled), "google_cloud_storage", or "aws_secrets_manager" for secure secret management | +| `cli_mode` | false | | +| `output_relevant_configurations` | false | | +| `output_run_details` | false | if true, append an agent run details section (model, tokens, time cost, AI calls) to generated PR comments | +| `output_run_cost` | false | if true, collect estimated LiteLLM API cost and include it inside the enabled run details section | +| `large_patch_policy` | "clip" | "clip", "skip" | +| `duplicate_prompt_examples` | false | | +| `persistent_inline_comments` | false | Persistent inline comments (issue #2037): when true, the GitHub, GitLab, and Azure DevOps providers fingerprint each inline comment, embed the fingerprint as an HTML marker, and skip re-posting suggestions already present on the PR/MR across runs. | +**seed** + +| Key | Default | Description | +| --- | --- | --- | +| `seed` | -1 | set positive value to fix the seed (and ensure temperature=0) | +| `temperature` | 0.2 | | +**ignore logic** + +| Key | Default | Description | +| --- | --- | --- | +| `ignore_pr_title` | ["^\\[Auto\\]", "^Auto"] | a list of regular expressions to match against the PR title to ignore the PR agent | +| `ignore_pr_target_branches` | [] | a list of regular expressions of target branches to ignore from PR agent when an PR is created | +| `ignore_pr_source_branches` | [] | a list of regular expressions of source branches to ignore from PR agent when an PR is created | +| `ignore_pr_labels` | [] | labels to ignore from PR agent when an PR is created | +| `ignore_pr_authors` | [] | authors to ignore from PR agent when an PR is created | +| `reaction_on_start` | "eyes" | added before the command runs | +| `reaction_on_success` | "" | added when the command finished successfully | +| `reaction_on_failure` | "" | added when the command failed | +| `ignore_repositories` | [] | a list of regular expressions of repository full names (e.g. "org/repo") to ignore from PR agent processing | +| `ignore_language_framework` | [] | a list of code-generation languages or frameworks (e.g. 'protobuf', 'go_gen') whose auto-generated source files will be excluded from analysis | +| `bot_user_indicators` | ["codium", "bot_", "bot-", "_bot", "-bot"] | Substring indicators used to skip bot users on webhook events. Currently consumed by the GitLab webhook (`is_bot_user`); other providers may adopt this list in future. The match is case-insensitive against the sender's display name. Overriding this setting REPLACES the default list — include the entries below in your override if you want to keep them (e.g. `["codium", "bot_", "bot-", "_bot", "-bot", "renovate"]`). | +| `restricted_mode` | false | when true, skip operations that require elevated permissions (e.g. pushing code to the repository) | +| `is_auto_command` | false | will be auto-set to true if the command is triggered by an automation | +| `propagate_tool_errors` | false | when true, a tool re-raises instead of swallowing an internal error, so a caller can tell a failed run from an empty one | +| `enable_ai_metadata` | false | will enable adding ai metadata | +| `add_user_to_requests` | false | send the current command and PR URL in the OpenAI-compatible "user" request field, for provider-side attribution of requests (e.g. OpenRouter "external_user") | +| `reasoning_effort` | "medium" | "none", "minimal", "low", "medium", "high", "xhigh", "max" | +**extended thinking for Claude reasoning models** + +| Key | Default | Description | +| --- | --- | --- | +| `enable_claude_extended_thinking` | false | Set to true to enable extended thinking feature | +| `extended_thinking_budget_tokens` | 2048 | | +| `extended_thinking_max_output_tokens` | 4096 | | +| `enable_claude_adaptive_thinking` | false | Adaptive thinking for Claude Opus 4.7/4.8 and Claude 5 models. When enabled, these models receive thinking={"type": "adaptive"} and an output_config effort. Do not add adaptive-only models to claude_extended_thinking_models_override; when both features are enabled, adaptive thinking takes precedence. | +| `claude_extended_thinking_models_override` | [] | Optional: override the built-in list of Claude models that receive the extended-thinking payload. When non-empty, this list fully replaces the built-in defaults (see CLAUDE_EXTENDED_THINKING_MODELS in pr_agent/algo/__init__.py). Leave empty to use the defaults. | +| `extract_issue_from_branch` | true | Extract issue number from PR source branch name (e.g. feature/1-auth-google -> issue #1). When true, branch-derived issue URLs are merged with tickets from the PR description for compliance. Set to false to restore description-only behaviour. Note: Branch-name extraction is GitHub-only for now; other providers planned for later. | +| `branch_issue_regex` | "" | Optional: custom regex with exactly one capturing group for the issue number (validated at runtime; falls back to default if missing). If empty, uses default pattern: first 1-6 digits at start of branch or after a slash, followed by hyphen or end (e.g. feature/1-test, 123-fix). GitHub only; other providers planned for later. | + + +## `[pr_reviewer]` — /review + +**enable/disable features** + +| Key | Default | Description | +| --- | --- | --- | +| `require_score_review` | false | | +| `require_tests_review` | true | | +| `require_estimate_effort_to_review` | true | | +| `require_can_be_split_review` | false | | +| `require_security_review` | true | | +| `require_estimate_contribution_time_cost` | false | | +| `require_todo_scan` | false | | +| `require_ticket_analysis_review` | true | | +| `require_risk_assessment` | false | ask the model for an overall risk level (low/medium/high) | +| `require_merge_recommendation` | false | ask the model for a merge recommendation (safe_to_merge/merge_with_caution/changes_required) | +| `require_priority_files` | false | ask the model which files a human should inspect first | +**general options** + +| Key | Default | Description | +| --- | --- | --- | +| `publish_output_no_suggestions` | true | Set to "false" if you only need the reviewer's remarks (not labels, not "security audit", etc.) and want to avoid noisy "No major issues detected" comments. | +| `publish_error_details` | false | Publish a deterministic, sanitized failure reason in manual review comments. No AI call is used. | +| `persistent_comment` | true | | +| `review_heading` | "PR Reviewer Guide" | Visible base heading for full and incremental review comments. Identity is tracked separately. | +| `persistent_finding_state` | true | Persist review finding state across complete review runs. | +| `inline_key_issues` | false | | +| `extra_instructions` | "" | | +| `num_max_findings` | 3 | | +| `final_update_message` | true | | +**review labels** + +| Key | Default | Description | +| --- | --- | --- | +| `enable_review_labels_security` | true | | +| `enable_review_labels_effort` | true | | +**specific configurations for incremental review (/review -i)** + +| Key | Default | Description | +| --- | --- | --- | +| `require_all_thresholds_for_incremental_review` | false | | +| `minimal_commits_for_incremental_review` | 0 | | +| `minimal_minutes_for_incremental_review` | 0 | | +| `enable_intro_text` | true | | +| `enable_help_text` | false | Determines whether to include help text in the PR review. | +| `enable_review_coverage_footer` | true | | +| `enable_large_pr_chunking` | false | large-diff chunking (opt-in). When the token budget leaves files out of the review, split the diff into chunks, review each chunk, and merge the per-chunk results into one review. | +| `max_number_of_calls` | 3 | maximum number of chunk review calls, used only when enable_large_pr_chunking is true | + + +## `[pr_description]` — /describe + +| Key | Default | Description | +| --- | --- | --- | +| `publish_labels` | false | | +| `add_original_user_description` | true | | +| `generate_ai_title` | false | | +| `extra_instructions` | "" | | +| `enable_pr_type` | true | | +| `enable_pr_description` | true | adds a section with an AI-generated summary of the PR changes | +| `final_update_message` | true | | +| `enable_help_text` | false | | +| `enable_help_comment` | false | | +| `enable_pr_diagram` | true | adds a section with a diagram of the PR changes | +| `pr_diagram_direction` | "adaptive" | 'adaptive', 'LR', 'TD'. 'adaptive' picks the direction from the shape of the diagram | +| `pr_diagram_direction_threshold` | 5 | with 'adaptive', a chain longer than this many nodes is drawn top-down | +**describe as comment** + +| Key | Default | Description | +| --- | --- | --- | +| `publish_description_as_comment` | false | | +| `publish_description_as_comment_persistent` | true | | +**changes walkthrough section** + +| Key | Default | Description | +| --- | --- | --- | +| `enable_semantic_files_types` | true | | +| `collapsible_file_list` | "adaptive" | true, false, 'adaptive' | +| `collapsible_file_list_threshold` | 6 | | +| `file_table_collapsible_open_by_default` | false | | +**markers** + +| Key | Default | Description | +| --- | --- | --- | +| `use_description_markers` | false | | +| `enable_large_pr_handling` | true | | +| `include_generated_by_header` | true | | +| `max_ai_calls` | 4 | | +| `async_ai_calls` | true | | + + +## `[pr_questions]` — /ask + +| Key | Default | Description | +| --- | --- | --- | +| `enable_help_text` | false | | +| `use_conversation_history` | true | | +| `ask_heading` | "Ask" | Set the plain-text heading for top-level /ask answers; provider-specific presentation is added automatically. | +| `resolve_threads` | false | Enable to let /ask_line resolve the review thread when the LLM judges the issue addressed. Note: also resolves threads started by human reviewers. | +| `extra_instructions` | "" | | + + +## `[pr_code_suggestions]` — /improve + +| Key | Default | Description | +| --- | --- | --- | +| `commitable_code_suggestions` | false | | +| `dual_publishing_score_threshold` | -1 | -1 to disable, [0-10] to set the threshold (>=) for publishing a code suggestion both in a table and as committable | +| `focus_only_on_problems` | true | | +| `extra_instructions` | "" | | +| `suggestions_heading` | "PR Code Suggestions" | Visible base heading for summary-table /improve comments. Identity is tracked separately. | +| `enable_help_text` | false | | +| `enable_chat_text` | false | | +| `persistent_comment` | true | | +| `max_history_len` | 4 | | +| `publish_output_no_suggestions` | true | | +| `enable_suggestions_coverage_footer` | true | show when failed analysis chunks make the suggestions incomplete | +**suggestions scoring** + +| Key | Default | Description | +| --- | --- | --- | +| `suggestions_score_threshold` | 0 | [0-10]\| recommend not to set this value above 8, since above it may clip highly relevant suggestions | +| `new_score_mechanism` | true | | +| `new_score_mechanism_th_high` | 9 | | +| `new_score_mechanism_th_medium` | 7 | | +**params for '/improve --extended' mode** + +| Key | Default | Description | +| --- | --- | --- | +| `num_code_suggestions_per_chunk` | 3 | | +| `max_suggestions_per_file` | 0 | Maximum suggestions retained per file after all chunks are merged; 0 disables the cap. | +| `max_number_of_calls` | 3 | | +| `parallel_calls` | true | | +| `decouple_hunks` | false | | +**self-review checkbox** + +| Key | Default | Description | +| --- | --- | --- | +| `demand_code_suggestions_self_review` | false | add a checkbox for the author to self-review the code suggestions | +| `code_suggestions_self_review_text` | "**Author self-review**: I have reviewed the PR code suggestions, and addressed the relevant ones." | | +| `approve_pr_on_self_review` | false | if true, the PR will be auto-approved after the author clicks on the self-review checkbox | +| `fold_suggestions_on_self_review` | true | if true, the code suggestions will be folded after the author clicks on the self-review checkbox | + + +## `[pr_add_docs]` — /add_docs + +| Key | Default | Description | +| --- | --- | --- | +| `extra_instructions` | "" | | +| `docs_style` | "Sphinx" | "Google Style with Args, Returns, Attributes...etc", "Numpy Style", "Sphinx Style", "PEP257", "reStructuredText" | + + +## `[pr_update_changelog]` — /update_changelog + +| Key | Default | Description | +| --- | --- | --- | +| `push_changelog_changes` | false | | +| `extra_instructions` | "" | | +| `add_pr_link` | true | | +| `skip_ci_on_push` | true | | + + +## `[pr_config]` — /config + +_This section only documents commented-out examples; see the [TOML source](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) for details._ + +## `[pr_help_docs]` + +| Key | Default | Description | +| --- | --- | --- | +| `repo_url` | "" | If not overwritten, will use the repo from where the context came from (issue or PR) | +| `repo_default_branch` | "main" | | +| `docs_path` | "docs" | | +| `exclude_root_readme` | false | | +| `supported_doc_exts` | [".md", ".mdx", ".rst"] | | +| `enable_help_text` | false | | + + +## `[github]` + +| Key | Default | Description | +| --- | --- | --- | +| `deployment_type` | "user" | The type of deployment to create. Valid values are 'app' or 'user'. | +| `ratelimit_retries` | 5 | | +| `seconds_between_requests` | 0 | seconds between API requests; 0 = no pacing (1.59 behaviour) | +| `seconds_between_writes` | 0 | seconds between write calls; 0 = no pacing (1.59 behaviour) | +| `api_retries` | 0 | max retries per request with backoff; 0 = no retries (1.59 behaviour) | +| `polling_request_timeout` | 10 | total seconds for comment-history fallback; positive values capped at 60 | +| `base_url` | "https://api.github.com" | | +| `try_fix_invalid_inline_comments` | true | | +| `ignore_bot_pr` | true | | +| `publish_as_check_run` | false | when true, publish review/description/improve output as GitHub Checks instead of PR comments | + + +## `[github_action_config]` + +_This section only documents commented-out examples; see the [TOML source](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) for details._ + +## `[github_app]` + +| Key | Default | Description | +| --- | --- | --- | +| `override_deployment_type` | true | these toggles allows running the github app from custom deployments | +**settings for "pull_request" event** + +| Key | Default | Description | +| --- | --- | --- | +| `handle_pr_actions` | ["opened", "reopened", "ready_for_review"] | | +| `feedback_on_draft_pr` | false | | +| `review_states` | ["changes_requested"] | A submitted GitHub review can optionally trigger these commands. The empty default preserves current behavior. | +| `review_author_types` | ["User"] | | +| `review_commands` | [] | | +| `webhook_delivery_deduplication` | false | Opt in to in-memory deduplication by X-GitHub-Delivery within each worker process. Active work stays protected; completed IDs expire after the deployment's push_trigger_pending_tasks_ttl (read at startup, 300 seconds by default). Failed or cancelled work can retry immediately. Repeated deliveries do not extend the TTL. Manual redeliveries are also suppressed during this window. State is cleared on restart and is not shared across workers or replicas. | +| `handle_push_trigger` | false | settings for "pull_request" event with "synchronize" action - used to detect and handle push triggers for new commits | +| `push_trigger_ignore_bot_commits` | true | | +| `push_trigger_ignore_merge_commits` | true | | +| `push_trigger_pending_tasks_backlog` | true | | +| `push_trigger_pending_tasks_ttl` | 300 | | +| `push_commands` | ["/describe", "/review"] | | + + +## `[gitlab]` + +| Key | Default | Description | +| --- | --- | --- | +| `url` | "https://gitlab.com" | | +| `expand_submodule_diffs` | false | | +| `feedback_on_draft_pr` | false | | +| `publish_review_as_thread` | false | Post the /review summary as a resolvable thread (discussion) instead of a plain note. | +| `publish_improve_as_thread` | false | Post the /improve suggestions comment as a resolvable thread (discussion) instead of a plain note. | +| `publish_code_suggestions_as_review` | false | When pr_code_suggestions.commitable_code_suggestions is true, queue each suggestion as a GitLab draft note and publish them all together in one batch (like GitLab's own "start a review" flow) instead of posting each as its own live discussion - and its own notification - as soon as it's created. | +| `resolve_outdated_inline_threads` | false | Resolve the bot's own inline threads that a later push left on an outdated diff version. | +| `handle_push_trigger` | false | | +| `push_commands` | ["/describe", "/review"] | | +| `handle_reviewer_assignment` | false | Auto-trigger commands when the bot is assigned as a reviewer on an MR | +| `reviewer_commands` | ["/review"] | | + + +## `[gitea]` + +| Key | Default | Description | +| --- | --- | --- | +| `url` | "https://gitea.com" | | +| `handle_push_trigger` | false | | +| `push_commands` | ["/describe", "/review"] | | + + +## `[bitbucket_app]` + +| Key | Default | Description | +| --- | --- | --- | +| `avoid_full_files` | false | | +| `request_timeout` | 30 | positive seconds for connection and response-read inactivity timeouts | + + +## `[local]` + +_This section only documents commented-out examples; see the [TOML source](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) for details._ + +## `[gerrit]` + +| Key | Default | Description | +| --- | --- | --- | +| `webhook_username` | "" | both are required: the gerrit webhook endpoint rejects every call until they are set | +| `webhook_password` | "" | | + + +## `[bitbucket_server]` + +| Key | Default | Description | +| --- | --- | --- | +| `url` | "" | URL to the BitBucket Server instance | + + +## `[jira]` + +| Key | Default | Description | +| --- | --- | --- | +| `jira_requirements_field` | "" | Custom field id holding acceptance criteria / requirements, mapped to the ticket "requirements" section. Instance-specific (e.g. "customfield_10127"); empty disables it. | +| `project_keys` | [] | Optional allowlist of Jira project keys, e.g. ["PROJ", "OPS"]. When non-empty, key-shaped text with another prefix ("SHA-256", "UTF-8", "ISO-8601") is dropped before any lookup, so it no longer costs an authenticated 404 each. Entries must be plain upper-case keys; a supplied list with no valid entry disables the lookup rather than widening it. Empty (default) looks up every key found. | + + +## `[litellm]` + +| Key | Default | Description | +| --- | --- | --- | +| `enable_callbacks` | false | | +| `success_callback` | [] | | +| `failure_callback` | [] | | +| `service_callback` | [] | | +| `turn_off_message_logging` | false | set true to keep prompt/response content (the whole PR diff) out of callback payloads; strongly recommended with the "otel" callback | +| `custom_llm_provider` | "" | Optional: forward a fixed custom_llm_provider to LiteLLM, so a raw hosted model id (e.g. "claude-sonnet-4-5") reaches the provider unchanged instead of being rewritten by LiteLLM's prefix inference. Empty = let LiteLLM infer the provider from the model name. | +| `force_streaming_custom_llm_provider` | "" | Force streaming when the request matches this provider AND its api_base contains one of the substrings below. Some OpenAI-compatible endpoints return a response that LiteLLM cannot normalize in non-streaming mode. Both must be set for the workaround to apply. | +| `force_streaming_api_base_substrings` | [] | | +| `callback_timeout_seconds` | 30 | max seconds to wait for pending litellm callbacks to flush before exiting | +| `cache_control_injection_points` | [] | Optional: enable Anthropic prompt caching via LiteLLM, e.g. [{location = "message", role = "system"}] (https://docs.litellm.ai/docs/tutorials/prompt_caching) | + + +## `[openrouter]` + +| Key | Default | Description | +| --- | --- | --- | +| `provider_only` | [] | restrict routing to these upstream providers only, a hard allowlist (e.g. ["z-ai"]); empty = OpenRouter default routing | +| `provider_order` | [] | preferred provider order; ignored when provider_only is set; empty = unset | +| `allow_fallbacks` | true | when provider_order is set, allow routing beyond the listed providers | +| `reasoning_effort` | "" | Invalid reasoning_effort values are warned about and treated as unset. Empty inherits config.reasoning_effort for models in SUPPORT_REASONING_EFFORT_MODELS. Valid values: "none", "minimal", "low", "medium", "high", "xhigh", "max". OpenRouter normalizes "max" to "xhigh" for LiteLLM/OpenRouter compatibility. Model-specific support varies; mandatory reasoning models reject "none". | +| `reasoning_max_tokens` | 0 | A positive value overrides global effort and non-none OpenRouter-specific efforts. Explicit openrouter.reasoning_effort = "none" keeps reasoning disabled, except on Grok 4.5/4.6: there "none" is clamped to the lowest supported effort first, so a positive budget wins over it. Some providers require max_tokens to be greater than the reasoning budget. | +| `max_tokens` | 0 | hard cap on completion tokens for the request; 0 = unset | + + +## `[model_routing]` — send a small pull request to a cheaper primary model (disabled by default) + +| Key | Default | Description | +| --- | --- | --- | +| `enable` | false | Rules are checked in order; the first one whose limits the pull request fits selects the primary model for the call, and config.fallback_models still apply after it. A pull request that fits no rule uses config.model. Size is measured after the [ignore] rules by diff hunks (max_hunks) and changed files (max_files), which every git provider reports and which do not depend on any model's tokenizer. Only calls that ask for the regular model are routed (/review, /improve, /generate_labels, /add_docs); tools that use model_weak are left alone. With Azure (openai.deployment_id set) a rule also needs its own deployment_id, or it is skipped. | +| `rules` | [] | e.g. [{ max_hunks = 3, model = "gpt-5.6-luna" }, { max_hunks = 15, max_files = 6, model = "gpt-5.6-terra" }] | + + +## `[otel]` + +| Key | Default | Description | +| --- | --- | --- | +| `is_enabled` | false | disabled by default; set to true to enable telemetry | +| `exporter_type` | "console" | "console", "otlp", "prometheus", or "none" | +| `service_name` | "pr-agent" | | +| `environment` | "development" | "development", "staging", "production", etc. | +| `otlp_timeout` | 3 | seconds; hard deadline per OTLP export call (incl. retries). Bounds CLI-exit/request stalls when the collector is unreachable; batches slower than this are dropped. | +| `otlp_protocol` | "http" | "http" (default; exporter ships with pr-agent) or "grpc" (requires the otel-grpc extra: pip install pr-agent[otel-grpc]) | +| `prometheus_multiproc_dir` | "/tmp/pr-agent-prometheus" | shared dir for the prometheus exporter's per-worker state files; required when exporter_type = "prometheus" (multiprocess gunicorn deployments); created 0700 and must stay owned by the run user (symlink or foreign-owner paths are refused) | +| `include_pr_url` | false | set to true to attach PR URLs to spans (may expose private repo names) | +| `include_error_details` | false | set to true to attach exception messages and rejected-command text to spans (may expose PR URLs, repo names, or other request content) | + + +## `[pr_similar_issue]` + +| Key | Default | Description | +| --- | --- | --- | +| `skip_comments` | false | | +| `force_update_dataset` | false | | +| `max_issues_to_scan` | 500 | | +| `vectordb` | "lancedb" | options: "pinecone", "lancedb", "qdrant" | + + +## `[pinecone]` + +| Key | Default | Description | +| --- | --- | --- | +| `cloud` | "aws" | Serverless index deployment for the modern SDK. `cloud` is one of "aws", "gcp" or "azure"; pick a `region` offered by that cloud. | +| `region` | "us-east-1" | | + + +## `[lancedb]` + +| Key | Default | Description | +| --- | --- | --- | +| `uri` | "./lancedb" | | + + +## `[qdrant]` + +_This section only documents commented-out examples; see the [TOML source](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) for details._ + +## `[skills]` + +| Key | Default | Description | +| --- | --- | --- | +| `enabled` | false | Agent skills (SKILL.md) support: discovers SKILL.md files from the configured filesystem paths and injects their content into review/improve/describe and top-level /ask prompts. Sibling *.md files in the skill directory tree (e.g. references/guide.md) are inlined alongside SKILL.md. PR-Agent supports text-only skills: scripts/ and assets/ subdirectories are skipped because PR-Agent uses a single-shot model call (no tool-use loop) and cannot execute scripts or load binary assets on demand. Skills that depend on script execution will not work here. See https://github.com/The-PR-Agent/pr-agent/issues/2384 | +| `paths` | [] | directories to scan recursively for "*/SKILL.md"; supports ~ and $VAR | +| `max_skills_tokens` | 8000 | token budget for the combined skills_context block | + + +## `[artifacts]` + +| Key | Default | Description | +| --- | --- | --- | +| `enable` | false | Enable artifact injection into tool prompts (off by default; auto-enabled when artifact_path input is set) | +| `artifact_path` | "" | File path to the artifact (relative to GITHUB_WORKSPACE, or absolute) | +| `artifact_instructions` | "" | Custom instructions appended after the artifact content (leave empty for a sensible default) | +| `artifact_label` | "" | Label shown to the AI — defaults to the filename when empty. | +| `target_tools` | ["pr_reviewer", "pr_description", "pr_code_suggestions"] | Which tools receive artifact context. | +| `max_artifact_size` | 50000 | Max artifact size in characters (content is truncated if exceeded) | + + +## `[asana]` + +| Key | Default | Description | +| --- | --- | --- | +| `request_timeout` | 10 | seconds allowed for each Asana task API request | + + +## `[azure_devops]` + +| Key | Default | Description | +| --- | --- | --- | +| `default_comment_status` | "closed" | | + + +## `[azure_devops_server]` + +| Key | Default | Description | +| --- | --- | --- | +| `agent_identity` | "" | empty: discover the identity from earlier agent comments on the PR | + + +## `[push_outputs]` — push tool outputs to external sinks without calling git-provider APIs (disabled by default) + +| Key | Default | Description | +| --- | --- | --- | +| `enable` | false | | +| `channels` | [] | any of: "stdout", "file", "webhook", "slack". Nothing is emitted until a channel is listed here | +| `file_path` | "pr-agent-outputs/reviews.jsonl" | used by the "file" channel | +| `webhook_url` | "" | used by the "webhook" channel: generic JSON POST target. Must be an absolute https:// URL | +| `slack_webhook_url` | "" | used by the "slack" channel: a Slack Incoming Webhook URL. Must be an absolute https:// URL | diff --git a/docs/docs/usage-guide/custom_ca_and_self_signed_certificates.md b/docs/docs/usage-guide/custom_ca_and_self_signed_certificates.md new file mode 100644 index 0000000000..cf2cd88f25 --- /dev/null +++ b/docs/docs/usage-guide/custom_ca_and_self_signed_certificates.md @@ -0,0 +1,46 @@ +## Custom CA and self-signed certificates + +When PR-Agent runs behind a corporate TLS-inspecting proxy or against servers using self-signed certificates, HTTPS calls to the LLM provider and git operations can fail with `certificate verify failed` errors. The fix requires two pieces: telling the HTTP clients which CA bundle to trust, and making LiteLLM use a transport that honours that bundle. + +### Environment variables for CA trust + +PR-Agent reads three environment variables when building the git clone environment, in the order below. The first one that is **set** wins, and if it points to a file that does not exist PR-Agent logs a warning and configures no bundle for git rather than falling through to the next variable: + +| Variable | Scope | Notes | +|---|---|---| +| `SSL_CERT_FILE` | git clone and Python `requests` | Preferred. Used by most Python SSL contexts. | +| `REQUESTS_CA_BUNDLE` | git clone and Python `requests` | Fallback. Honoured by the `requests` library directly. | +| `GIT_SSL_CAINFO` | git clone only | Last resort. Used by git when the other two are absent. | + +If more than one variable is set and points to a different file, PR-Agent logs a warning and picks one according to the precedence above. All three variables should point to the same PEM file in practice. + +Example (runner environment or shell profile): + +```bash +export SSL_CERT_FILE=/etc/ssl/certs/corporate-ca-bundle.crt +``` + +These variables are read from the process environment, so they must be exported in the runner or shell profile; a `.secrets.toml` `[env]` section does **not** reach `os.environ` and cannot set them. + +### LiteLLM transport fallback + +By default LiteLLM uses aiohttp for HTTP calls, and aiohttp does not honour the standard CA environment variables. Setting `litellm.disable_aiohttp` makes LiteLLM fall back to httpx, which reads `SSL_CERT_FILE` (and `SSL_CERT_DIR`). httpx does not read `REQUESTS_CA_BUNDLE`, so the LLM call needs `SSL_CERT_FILE` specifically. + +```toml +[litellm] +disable_aiohttp = true +``` + +This setting is read once when the LiteLLM handler initialises, so it must be present before the process starts (runner environment or `.secrets.toml`, not a runtime override). + +### Scope of `gitlab.ssl_verify` + +The `gitlab.ssl_verify` setting (or `gitlab__SSL_VERIFY` environment variable) is passed only to the python-gitlab client that talks to the GitLab API. It does **not** affect the LLM call, git clone operations, or any other HTTPS client in the process. If you are seeing certificate errors from LiteLLM or git clone, use the environment variables above instead. + +### Putting it all together + +For a runner behind a corporate proxy with a custom CA: + +1. Export `SSL_CERT_FILE` pointing to your CA bundle. `REQUESTS_CA_BUNDLE` and `GIT_SSL_CAINFO` cover git clone only. +2. Set `litellm.disable_aiohttp = true` in your configuration. +3. If you also use the GitLab API, keep `gitlab.ssl_verify` pointed at the same bundle for the python-gitlab client. diff --git a/docs/docs/usage-guide/index.md b/docs/docs/usage-guide/index.md index bb08ee6f1c..3f37280265 100644 --- a/docs/docs/usage-guide/index.md +++ b/docs/docs/usage-guide/index.md @@ -5,6 +5,7 @@ It includes information on how to adjust PR-Agent configurations, define which t - [Introduction](./introduction.md) - [Configuration File](./configuration_options.md) +- [Configuration Reference](./configuration_reference.md) - [Usage and Automation](./automations_and_usage.md) - [Local Repo (CLI)](./automations_and_usage.md#local-repo-cli) - [Online Usage](./automations_and_usage.md#online-usage) @@ -15,10 +16,10 @@ It includes information on how to adjust PR-Agent configurations, define which t - [BitBucket App](./automations_and_usage.md#bitbucket-app) - [Azure DevOps Provider](./automations_and_usage.md#azure-devops-provider) - [Managing Mail Notifications](./mail_notifications.md) +- [Push Outputs](./push_outputs.md) - [Changing a Model](./changing_a_model.md) - [Additional Configurations](./additional_configurations.md) - [Ignoring files from analysis](./additional_configurations.md#ignoring-files-from-analysis) - [Extra instructions](./additional_configurations.md#extra-instructions) - - [Working with large PRs](./additional_configurations.md#working-with-large-prs) - [Changing a model](./changing_a_model.md) - [FAQ](../faq/index.md) diff --git a/docs/docs/usage-guide/mail_notifications.md b/docs/docs/usage-guide/mail_notifications.md index ca6933be16..cfa6f6a32e 100644 --- a/docs/docs/usage-guide/mail_notifications.md +++ b/docs/docs/usage-guide/mail_notifications.md @@ -2,11 +2,11 @@ Unfortunately, it is not possible in GitHub to disable mail notifications from a specific user. If you are subscribed to notifications for a repo with PR-Agent, we recommend turning off notifications for PR comments, to avoid lengthy emails: -![notifications](https://codium.ai/images/pr_agent/notifications.png){width=512} +![notifications](../assets/notifications.png){width=512} As an alternative, you can filter in your mail provider the notifications specifically from the PR-Agent bot, [see how](https://www.quora.com/How-can-you-filter-emails-for-specific-people-in-Gmail#:~:text=On%20the%20Filters%20and%20Blocked,the%20body%20of%20the%20email). -![filter_mail_notifications](https://codium.ai/images/pr_agent/filter_mail_notifications.png){width=512} +![filter_mail_notifications](../assets/filter_mail_notifications.png){width=512} Another option to reduce the mail overload, yet still receive notifications on PR-Agent tools, is to disable the help collapsible section in PR-Agent bot comments. This can done by setting `enable_help_text=false` for the relevant tool in the configuration file. diff --git a/docs/docs/usage-guide/plain_diff_mode.md b/docs/docs/usage-guide/plain_diff_mode.md index d754e771a9..560dd59c68 100644 --- a/docs/docs/usage-guide/plain_diff_mode.md +++ b/docs/docs/usage-guide/plain_diff_mode.md @@ -26,20 +26,28 @@ python -m pr_agent.cli --diff-file changes.diff --output review.md review |---|---| | `--stdin` | Read a unified diff from stdin | | `--diff-file ` | Read a unified diff from a file | -| `--output ` | Write the result to a file in addition to stdout | -| `--json-output ` | Write the parsed review and token usage to a JSON file (`review` command) | +| `--output ` | Write Markdown produced by a compatible Plain Diff command to a file in addition to stdout | +| `--json-output ` | Write the parsed review and token usage to a JSON file (`review` or `review_pr` only) | `--stdin` and `--diff-file` are mutually exclusive. At least one must be provided to enter plain-diff mode; omitting both falls back to the normal `--pr_url` flow. +Output options must appear before the command. For example, use +`--stdin --output result.md review`, not `--stdin review --output result.md`. ### Supported commands -`review`, `improve`, `describe`, and `ask` are supported. Because there is no hosting +`review`, `improve`, `describe`, and `ask` are supported, along with their +`review_pr`, `improve_code`, `describe_pr`, and `ask_question` aliases. Because there is no hosting platform to push to, `improve` renders its code suggestions as a single markdown document to stdout (and to `--output`, if given) instead of as committable inline suggestions. Commands that require live platform interaction (such as `update_changelog` or `similar_issue`) are not meaningful in this mode. +Existing command variants that publish Markdown without additional platform +state (`auto_review`, `config`, `settings`, and `help`) also honor `--output`. +`answer` does not: it requires issue-comment history, which is unavailable when +the input is a standalone diff. + ## How it works 1. **Diff parsing** — the unified diff is parsed locally into per-file patch objects. @@ -55,9 +63,10 @@ suggestions. Commands that require live platform interaction (such as for that file. The review still runs; it simply has less context. 4. **Output** — the result is written to stdout. If `--output ` is given, it is - also written to that file (UTF-8, overwritten on each run). For `review`, if - `--json-output ` is given, the parsed review plus a `usage` object with the - run's accumulated token counts is also written to that file as JSON. + also written to that file (UTF-8, overwritten on each run). For `review` or + `review_pr`, if `--json-output ` is given, the parsed review plus a + `usage` object with the run's accumulated token counts is also written to that + file as JSON. No platform token, no PR URL, and no internet access are required for the diff processing step itself. An LLM API key is still needed unless you configure a local model. diff --git a/docs/docs/usage-guide/push_outputs.md b/docs/docs/usage-guide/push_outputs.md new file mode 100644 index 0000000000..7fabba0dbc --- /dev/null +++ b/docs/docs/usage-guide/push_outputs.md @@ -0,0 +1,90 @@ +# Push outputs to external sinks + +The `[push_outputs]` feature routes finished tool output to external sinks — stdout, a JSONL file, a +generic webhook, or Slack — without calling git-provider APIs. It is disabled by default, and is +additive to normal publishing: when a tool finishes, the same result that is posted as a PR comment +is also emitted to the configured sinks. + +## What gets pushed + +Each finished tool run emits one record. The tools that currently emit are: + +| Tool | `type` in the record | +|---|---| +| `/review` | `review` | +| `/describe` | `describe` | +| `/improve` | `improve` | + +A record is a JSON object: + +```json +{ + "type": "review", + "timestamp": "2026-09-08T14:03:21+00:00", + "payload": {}, + "markdown": "# PR Reviewer Guide ..." +} +``` + +- `type` — which tool produced the output +- `timestamp` — the run completion time, UTC, ISO-8601 +- `payload` — the structured tool result +- `markdown` — the rendered comment text; present when the tool produces one + +## Configuration + +The defaults are defined at the end of the +[configuration file](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml): + +```toml +[push_outputs] +enable = false +channels = [] # any of: "stdout", "file", "webhook", "slack" +file_path = "pr-agent-outputs/reviews.jsonl" +webhook_url = "" # must be an absolute https:// URL +slack_webhook_url = "" # Slack Incoming Webhook; must be an absolute https:// URL +``` + +- `enable` — master switch (default `false`). When `false`, nothing is emitted. +- `channels` — which sinks to use. Nothing is emitted until at least one channel is listed here. +- `file_path` — the file the `file` channel appends to. +- `webhook_url` — the endpoint the `webhook` channel POSTs the generic record to. +- `slack_webhook_url` — a Slack Incoming Webhook URL that the `slack` channel posts a `{"text": ...}` payload to. + +!!! danger "Host-only configuration" + The whole `[push_outputs]` section is **host-only**. A repository cannot set these keys: + keys supplied through a repo's local `.pr_agent.toml` are dropped, and CLI arguments + (`--push_outputs.webhook_url=...`, `--push_outputs={...}`) are blocked. This prevents a + pull request from redirecting review output to an attacker-controlled host, reaching + internal endpoints, or appending to arbitrary host files. Configure these values in the + PR-Agent host's own settings. + +### URL requirements + +`webhook_url` and `slack_webhook_url` must be absolute `https://` URLs with a host. Any other +value (for example a plain `http://` URL or a bare path) is ignored with a warning. Requiring +HTTPS keeps review text, which can quote private code, off plaintext transports. The host is +intentionally not restricted, so self-hosted collectors and Slack-compatible endpoints +(Mattermost, Rocket.Chat) are legitimate targets. + +Warnings log the setting name — never the URL value — because a webhook or Slack URL is itself a +credential. + +## Channels + +| Channel | Behaviour | +|---|---| +| `stdout` | Prints one JSON line (the record) to stdout. | +| `file` | Appends one JSON line per run (JSONL) to `file_path`, creating parent directories as needed. | +| `webhook` | POSTs the generic record as JSON to `webhook_url` (5-second timeout, redirects not followed). | +| `slack` | POSTs `{"text": ...}` to a Slack Incoming Webhook; the text is the markdown, or the payload JSON when the tool produces no markdown. | + +Local channels (`stdout`, `file`) run before network channels (`webhook`, `slack`), and network +posts never follow redirects. Each configured destination is attempted independently, so one failure +does not prevent later destinations from receiving the output. + +## Error handling + +Failures are non-fatal: `push_outputs` never raises, so a sink outage does not break the review +flow. Exceptions and non-2xx HTTP responses are logged with the destination and only the exception +type or status code, since request error messages can embed the (secret-bearing) URL. diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 98774f1352..cc501eb20b 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -15,15 +15,19 @@ nav: - BitBucket Integration: 'installation/bitbucket.md' - Azure DevOps Integration: 'installation/azure.md' - Gitea Integration: 'installation/gitea.md' + - MOSAICO A2A Server: 'installation/mosaico_server.md' - Usage Guide: - 'usage-guide/index.md' - Introduction: 'usage-guide/introduction.md' - Configuration File: 'usage-guide/configuration_options.md' + - Configuration Reference: 'usage-guide/configuration_reference.md' - Usage and Automation: 'usage-guide/automations_and_usage.md' - Managing Mail Notifications: 'usage-guide/mail_notifications.md' + - Push Outputs: 'usage-guide/push_outputs.md' - Changing a Model: 'usage-guide/changing_a_model.md' - Extending PR-Agent: 'usage-guide/extending_pr_agent.md' - Additional Configurations: 'usage-guide/additional_configurations.md' + - Custom CA and Self-Signed Certificates: 'usage-guide/custom_ca_and_self_signed_certificates.md' - Plain-Diff Mode: 'usage-guide/plain_diff_mode.md' - Local Git Provider: 'usage-guide/local_git_provider.md' - Frequently Asked Questions: 'faq/index.md' diff --git a/pr_agent/agent/pr_agent.py b/pr_agent/agent/pr_agent.py index 2414055b6c..f5097d702a 100644 --- a/pr_agent/agent/pr_agent.py +++ b/pr_agent/agent/pr_agent.py @@ -178,7 +178,7 @@ def prepare_command(command: str) -> list[str]: class PRAgent: def __init__(self, ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler): - self.ai_handler = ai_handler # will be initialized in run_action + self.ai_handler = ai_handler # handler factory passed to each tool when it is instantiated async def _handle_request(self, pr_url, request, notify=None) -> bool: # Exceptions raised inside are caught below, but a BaseException (e.g. the diff --git a/pr_agent/algo/__init__.py b/pr_agent/algo/__init__.py index 01ed447e95..b326d217c3 100644 --- a/pr_agent/algo/__init__.py +++ b/pr_agent/algo/__init__.py @@ -269,6 +269,12 @@ def _generate_claude_registries(families=None): ) +# MAX_TOKENS holds only entries that deliberately deviate from (or are absent +# from) LiteLLM's model registry. Exact LiteLLM duplicates were removed because +# get_max_tokens() already falls back to litellm.get_model_info() with the same +# value; generator-expanded Claude families (see _CLAUDE_MODEL_FAMILIES) remain +# because they also drive the no-temperature / extended-thinking registries and +# their 1M-context handling is a separate, deliberate judgement (issue #3196). MAX_TOKENS = { 'text-embedding-ada-002': 8000, 'gpt-3.5-turbo': 16000, @@ -280,25 +286,8 @@ def _generate_claude_registries(families=None): 'gpt-4': 8000, 'gpt-4-0613': 8000, 'gpt-4-32k': 32000, - 'gpt-4-1106-preview': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4-0125-preview': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4o': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4o-2024-05-13': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4-turbo-preview': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4-turbo-2024-04-09': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4-turbo': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4o-mini': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4o-mini-2024-07-18': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4o-2024-08-06': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4o-2024-11-20': 128000, # 128K, but may be limited by config.max_model_tokens 'gpt-4.5-preview': 128000, # 128K, but may be limited by config.max_model_tokens 'gpt-4.5-preview-2025-02-27': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-4.1': 1047576, - 'gpt-4.1-2025-04-14': 1047576, - 'gpt-4.1-mini': 1047576, - 'gpt-4.1-mini-2025-04-14': 1047576, - 'gpt-4.1-nano': 1047576, - 'gpt-4.1-nano-2025-04-14': 1047576, 'gpt-5-nano': 200000, # 200K, but may be limited by config.max_model_tokens 'gpt-5-mini': 200000, # 200K, but may be limited by config.max_model_tokens 'gpt-5': 200000, @@ -310,7 +299,6 @@ def _generate_claude_registries(families=None): 'gpt-5.1-codex-mini': 200000, 'gpt-5.2': 400000, # 400K, but may be limited by config.max_model_tokens 'gpt-5.2-2025-12-11': 400000, # 400K, but may be limited by config.max_model_tokens - 'gpt-5.2-chat-latest': 128000, # 128K, but may be limited by config.max_model_tokens 'gpt-5.2-codex': 400000, # 400K, but may be limited by config.max_model_tokens 'gpt-5.3-codex': 400000, # 400K, but may be limited by config.max_model_tokens 'gpt-5.3-chat': 128000, # 128K, but may be limited by config.max_model_tokens @@ -320,8 +308,6 @@ def _generate_claude_registries(families=None): 'gpt-5.4-mini-2026-03-17': 400000, # 400K, but may be limited by config.max_model_tokens 'gpt-5.4-nano': 400000, # 400K, but may be limited by config.max_model_tokens 'gpt-5.4-nano-2026-03-17': 400000, # 400K, but may be limited by config.max_model_tokens - 'gpt-5.5': 1050000, # 1.05M, but may be limited by config.max_model_tokens - 'gpt-5.5-2026-04-23': 1050000, # 1.05M, but may be limited by config.max_model_tokens 'gpt-5.6': 1050000, # 1.05M, but may be limited by config.max_model_tokens 'gpt-5.6-sol': 1050000, # 1.05M, but may be limited by config.max_model_tokens 'gpt-5.6-terra': 1050000, # 1.05M, but may be limited by config.max_model_tokens @@ -335,17 +321,10 @@ def _generate_claude_registries(families=None): 'o1': 204800, # 200K, but may be limited by config.max_model_tokens 'o3-mini': 204800, # 200K, but may be limited by config.max_model_tokens 'o3-mini-2025-01-31': 204800, # 200K, but may be limited by config.max_model_tokens - 'o3': 200000, # 200K, but may be limited by config.max_model_tokens - 'o3-2025-04-16': 200000, # 200K, but may be limited by config.max_model_tokens - 'o4-mini': 200000, # 200K, but may be limited by config.max_model_tokens - 'o4-mini-2025-04-16': 200000, # 200K, but may be limited by config.max_model_tokens 'claude-instant-1': 100000, 'claude-2': 100000, - 'command-nightly': 4096, 'deepseek/deepseek-chat': 128000, # 128K, but may be limited by config.max_model_tokens 'deepseek/deepseek-reasoner': 64000, # 64K, but may be limited by config.max_model_tokens - 'deepseek/deepseek-v4-pro': 1000000, # 1M, but may be limited by config.max_model_tokens - 'deepseek/deepseek-v4-flash': 1000000, # 1M, but may be limited by config.max_model_tokens 'zai/glm-5.2': 200000, # 200K, matching the Z.AI GLM-5/5.1 lineage, but may be limited by config.max_model_tokens 'moonshot/kimi-k3': 262144, # 256K, matching the Moonshot Kimi-k2.5/k2.6 lineage, but may be limited by config.max_model_tokens 'openai/qwq-plus': 131072, # 131K context length, but may be limited by config.max_model_tokens @@ -362,8 +341,6 @@ def _generate_claude_registries(families=None): 'vertex_ai/claude-3-5-haiku@20241022': 100000, 'vertex_ai/claude-3-sonnet@20240229': 100000, 'vertex_ai/claude-3-opus@20240229': 100000, - 'vertex_ai/claude-opus-4@20250514': 200000, - 'vertex_ai/claude-opus-4-1@20250805': 200000, 'vertex_ai/claude-3-5-sonnet@20240620': 100000, 'vertex_ai/claude-3-5-sonnet-v2@20241022': 100000, 'vertex_ai/claude-sonnet-4@20250514': 200000, @@ -372,78 +349,45 @@ def _generate_claude_registries(families=None): 'vertex_ai/gemini-2.5-pro-preview-03-25': 1048576, 'vertex_ai/gemini-2.5-pro-preview-05-06': 1048576, 'vertex_ai/gemini-2.5-pro-preview-06-05': 1048576, - 'vertex_ai/gemini-2.5-pro': 1048576, 'vertex_ai/gemini-1.5-flash': 1048576, - 'vertex_ai/gemini-2.0-flash': 1048576, 'vertex_ai/gemini-2.5-flash-preview-04-17': 1048576, 'vertex_ai/gemini-2.5-flash-preview-05-20': 1048576, - 'vertex_ai/gemini-2.5-flash': 1048576, - 'vertex_ai/gemini-3-flash-preview': 1048576, - 'vertex_ai/gemini-3-pro-preview': 1048576, 'vertex_ai/gemini-3.1-flash': 1048576, 'vertex_ai/gemini-3.1-pro': 1048576, - 'vertex_ai/gemini-3.1-flash-lite-preview': 1048576, - 'vertex_ai/gemini-3.1-pro-preview': 1048576, - 'vertex_ai/gemini-3.5-flash': 1048576, - 'vertex_ai/gemini-3.5-flash-lite': 1048576, 'vertex_ai/gemini-3.5-pro': 1048576, - 'vertex_ai/gemini-3.6-flash': 1048576, - 'vertex_ai/gemini-3.7-flash': 1048576, - 'vertex_ai/gemini-3.8-flash': 1048576, + 'vertex_ai/gemini-3.8-flash': 1048576, # kept pinned: absent from LiteLLM's bundled cost map 'vertex_ai/gemma2': 8200, 'gemini/gemini-1.5-pro': 1048576, 'gemini/gemini-1.5-flash': 1048576, - 'gemini/gemini-2.0-flash': 1048576, 'gemini/gemini-2.5-flash-preview-04-17': 1048576, 'gemini/gemini-2.5-flash-preview-05-20': 1048576, - 'gemini/gemini-2.5-flash': 1048576, 'gemini/gemini-2.5-pro-preview-03-25': 1048576, 'gemini/gemini-2.5-pro-preview-05-06': 1048576, 'gemini/gemini-2.5-pro-preview-06-05': 1048576, - 'gemini/gemini-2.5-pro': 1048576, - 'gemini/gemini-3-flash-preview': 1048576, - 'gemini/gemini-3-pro-preview': 1048576, 'gemini/gemini-3.1-flash': 1048576, 'gemini/gemini-3.1-pro': 1048576, - 'gemini/gemini-3.1-flash-lite-preview': 1048576, - 'gemini/gemini-3.1-pro-preview': 1048576, - 'gemini/gemini-3.5-flash': 1048576, - 'gemini/gemini-3.5-flash-lite': 1048576, 'gemini/gemini-3.5-pro': 1048576, - 'gemini/gemini-3.6-flash': 1048576, - 'gemini/gemini-3.7-flash': 1048576, - 'gemini/gemini-3.8-flash': 1048576, + 'gemini/gemini-3.8-flash': 1048576, # kept pinned: absent from LiteLLM's bundled cost map 'codechat-bison': 6144, 'codechat-bison-32k': 32000, # -- Anthropic Claude -------------------------------------------------- - 'anthropic.claude-instant-v1': 100000, - 'anthropic.claude-v1': 100000, 'anthropic.claude-v2': 100000, 'anthropic/claude-3-opus-20240229': 100000, - 'anthropic/claude-opus-4-20250514': 200000, - 'anthropic/claude-opus-4-1-20250805': 200000, 'anthropic/claude-3-5-sonnet-20240620': 100000, 'anthropic/claude-3-5-sonnet-20241022': 100000, 'anthropic/claude-sonnet-4-20250514': 200000, # -- Bare Claude ------------------------------------------------------- - 'claude-opus-4-1-20250805': 200000, # -- Haiku ------------------------------------------------------------- 'anthropic/claude-3-5-haiku-20241022': 100000, # -- Bedrock Claude ---------------------------------------------------- - 'bedrock/anthropic.claude-instant-v1': 100000, 'bedrock/anthropic.claude-v2': 100000, - 'bedrock/anthropic.claude-v2:1': 100000, 'bedrock/anthropic.claude-3-sonnet-20240229-v1:0': 100000, - 'bedrock/anthropic.claude-opus-4-20250514-v1:0': 200000, - 'bedrock/anthropic.claude-opus-4-1-20250805-v1:0': 200000, 'bedrock/anthropic.claude-3-haiku-20240307-v1:0': 100000, 'bedrock/anthropic.claude-3-5-haiku-20241022-v1:0': 100000, 'bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0': 100000, 'bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0': 100000, 'bedrock/anthropic.claude-sonnet-4-20250514-v1:0': 200000, # -- Bedrock Claude (cross-region) ------------------------------------- - "bedrock/us.anthropic.claude-opus-4-20250514-v1:0": 200000, - "bedrock/us.anthropic.claude-opus-4-1-20250805-v1:0": 200000, "bedrock/us.anthropic.claude-3-5-sonnet-20241022-v2:0": 100000, "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0": 200000, "bedrock/global.anthropic.claude-sonnet-4-20250514-v1:0": 200000, @@ -451,17 +395,10 @@ def _generate_claude_registries(families=None): "bedrock/apac.anthropic.claude-sonnet-4-20250514-v1:0": 200000, 'claude-3-5-sonnet': 100000, # -- Non-Claude models ------------------------------------------------- - 'bedrock/us.meta.llama4-scout-17b-instruct-v1:0': 128000, - 'bedrock/us.meta.llama4-maverick-17b-instruct-v1:0': 128000, "bedrock_mantle/xai.grok-4.3": 1000000, # 1M context, but may be limited by config.max_model_tokens - 'groq/openai/gpt-oss-120b': 131072, - 'groq/openai/gpt-oss-20b': 131072, - 'groq/qwen/qwen3-32b': 131000, 'dashscope/qwen3.8-max': 1000000, # 1M, qwen3.8-max is the actual DashScope model id (context_window 1M per QwenCode metadata), but may be limited by config.max_model_tokens 'groq/moonshotai/kimi-k2-instruct': 131072, 'groq/deepseek-r1-distill-llama-70b': 128000, - 'groq/meta-llama/llama-4-maverick-17b-128e-instruct': 131072, - 'groq/meta-llama/llama-4-scout-17b-16e-instruct': 131072, 'groq/llama-3.3-70b-versatile': 128000, 'groq/llama-3.1-8b-instant': 128000, 'sambanova/MiniMax-M3': 192000, @@ -473,20 +410,10 @@ def _generate_claude_registries(families=None): 'xai/grok-2': 131072, 'xai/grok-2-1212': 131072, 'xai/grok-2-latest': 131072, - 'xai/grok-3': 131072, - 'xai/grok-3-beta': 131072, 'xai/grok-3-fast': 131072, - 'xai/grok-3-fast-beta': 131072, - 'xai/grok-3-mini': 131072, - 'xai/grok-3-mini-beta': 131072, - 'xai/grok-3-mini-fast': 131072, - 'xai/grok-3-mini-fast-beta': 131072, - "xai/grok-4.5": 500000, # 500K context, but may be limited by config.max_model_tokens - "xai/grok-4.5-latest": 500000, - "xai/grok-build-latest": 500000, - "xai/grok-4.6": 500000, # 500K context, but may be limited by config.max_model_tokens - "openrouter/x-ai/grok-4.5": 500000, - "openrouter/x-ai/grok-4.6": 500000, + "xai/grok-build-latest": 500000, # kept pinned: absent from LiteLLM's bundled cost map + "openrouter/x-ai/grok-4.5": 500000, # kept pinned: absent from LiteLLM's bundled cost map + "openrouter/x-ai/grok-4.6": 500000, # kept pinned: absent from LiteLLM's bundled cost map 'ollama/llama3': 4096, 'watsonx/meta-llama/llama-3-8b-instruct': 4096, "watsonx/meta-llama/llama-3-70b-instruct": 4096, @@ -499,16 +426,11 @@ def _generate_claude_registries(families=None): "deepinfra/deepseek-ai/DeepSeek-R1": 128000, "mistral/mistral-small-latest": 8191, "mistral/mistral-medium-latest": 8191, - "mistral/mistral-large-2407": 128000, "mistral/mistral-large-latest": 128000, "mistral/open-mistral-7b": 8191, "mistral/open-mixtral-8x7b": 8191, "mistral/open-mixtral-8x22b": 8191, "mistral/codestral-latest": 8191, - "mistral/open-mistral-nemo": 128000, - "mistral/open-mistral-nemo-2407": 128000, - "mistral/open-codestral-mamba": 256000, - "mistral/codestral-mamba-latest": 256000, "codestral/codestral-latest": 8191, "codestral/codestral-2405": 8191, 'xiaomi_mimo/mimo-v2.5': 1048576, # 1M, matching the LiteLLM registry for mimo-v2.5, xiaomi_mimo/ is the native LiteLLM Xiaomi provider, but may be limited by config.max_model_tokens diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index 3d22f0dbf6..46624faf49 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -1,8 +1,18 @@ import asyncio +import configparser import contextlib +import copy +import hashlib +import inspect import json import os import re +import shutil +import stat +from contextvars import ContextVar +from functools import lru_cache, wraps +from types import FunctionType, SimpleNamespace +from urllib.parse import urlparse import httpx import litellm @@ -11,6 +21,27 @@ from litellm import acompletion from tenacity import retry, retry_if_exception, stop_after_attempt +try: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo +except ImportError: + AnthropicModelInfo = None + +try: + from litellm.llms.openai_like.json_loader import JSONProviderRegistry +except ImportError: + JSONProviderRegistry = None + +try: + from litellm.utils import _get_model_info_helper +except ImportError: + _get_model_info_helper = None + +try: + from litellm.llms.bedrock_mantle.common_utils import MANTLE_HOST_RE, BedrockMantleAuthMixin +except ImportError: + BedrockMantleAuthMixin = None + MANTLE_HOST_RE = None + from pr_agent.algo import ( CLAUDE_EXTENDED_THINKING_MODELS, GROK_REASONING_EFFORT_LEVELS, @@ -22,6 +53,7 @@ ) from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.litellm_helpers import ( + _get_azure_ad_credential, _get_azure_ad_token, _handle_streaming_response, _process_litellm_extra_body, @@ -34,8 +66,1822 @@ from pr_agent.log import get_logger MODEL_RETRIES = 2 -DUMMY_LITELLM_API_KEY = "dummy_key" # placeholder set when no OpenAI key is configured +DUMMY_LITELLM_API_KEY = "dummy_key" # request-local guard against LiteLLM's process-wide key fallbacks _IMAGE_HEAD_TIMEOUT_SECONDS = 5 +OPENAI_DEFAULT_API_BASE = "https://api.openai.com/v1" + +PROVIDER_SETTING_PATHS = { + "anthropic": {"api_key": "ANTHROPIC.KEY"}, + "codestral": {"api_key": "CODESTRAL.KEY"}, + "cohere": {"api_key": "COHERE.KEY"}, + "cohere_chat": {"api_key": "COHERE.KEY"}, + "dashscope": {"api_key": "DASHSCOPE.KEY"}, + "databricks": {"api_key": "DATABRICKS.API_KEY", "api_base": "DATABRICKS.API_BASE"}, + "deepinfra": {"api_key": "DEEPINFRA.KEY"}, + "deepseek": {"api_key": "DEEPSEEK.KEY"}, + "gemini": {"api_key": "GOOGLE_AI_STUDIO.GEMINI_API_KEY"}, + "groq": {"api_key": "GROQ.KEY"}, + "huggingface": {"api_key": "HUGGINGFACE.KEY", "api_base": "HUGGINGFACE.API_BASE"}, + "mistral": {"api_key": "MISTRAL.KEY"}, + "moonshot": {"api_key": "MOONSHOT.KEY", "api_base": "MOONSHOT.API_BASE"}, + "ollama": {"api_key": "OLLAMA.API_KEY", "api_base": "OLLAMA.API_BASE"}, + "openrouter": {"api_key": "OPENROUTER.KEY", "api_base": "OPENROUTER.API_BASE"}, + "replicate": {"api_key": "REPLICATE.KEY"}, + "sambanova": {"api_key": "SAMBANOVA.KEY"}, + "text-completion-codestral": {"api_key": "CODESTRAL.KEY"}, + "xai": {"api_key": "XAI.KEY"}, + "xiaomi_mimo": {"api_key": "XIAOMI_MIMO.KEY"}, + "zai": {"api_key": "ZAI.KEY"}, +} + +PROVIDER_SETTING_ALIASES = { + "aiohttp_openai": "openai", + "anthropic_text": "anthropic", + "azure_text": "azure", + "ollama_chat": "ollama", + "text-completion-inception": "inception", + "text-completion-openai": "openai", + "vertex_ai_beta": "vertex_ai", +} + +# Keep chat-completion endpoint aliases in the same precedence order as LiteLLM 1.100.0. +# Native completion masks BASETEN_API_BASE, MISTRAL_API_BASE and ARK_API_BASE; +# VERTEX_API_BASE belongs to embedding, not chat. Do not promote them to explicit routing. +PROVIDER_API_BASE_ENV_VARS = { + "a2a": ("A2A_API_BASE",), + "ai21": ("AI21_API_BASE",), + "ai21_chat": ("AI21_API_BASE",), + "aiml": ("AIML_API_BASE",), + "aleph_alpha": ("ALEPH_ALPHA_API_BASE",), + "amazon_nova": ("AMAZON_NOVA_API_BASE",), + "anthropic": ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"), + "anyscale": ("ANYSCALE_API_BASE",), + "azure_ai": ("AZURE_AI_API_BASE",), + "bedrock_mantle": ("BEDROCK_MANTLE_API_BASE",), + "cerebras": ("CEREBRAS_API_BASE",), + "chatgpt": ("CHATGPT_API_BASE", "OPENAI_CHATGPT_API_BASE"), + "cloudflare": ("CLOUDFLARE_API_BASE",), + "codestral": ("CODESTRAL_API_BASE",), + "text-completion-codestral": ("CODESTRAL_API_BASE",), + "cohere": ("COHERE_API_BASE",), + "cohere_chat": ("COHERE_API_BASE",), + "cometapi": ("COMETAPI_API_BASE",), + "dashscope": ("DASHSCOPE_API_BASE",), + "databricks": ("DATABRICKS_API_BASE",), + "datarobot": ("DATAROBOT_ENDPOINT",), + "deepinfra": ("DEEPINFRA_API_BASE",), + "deepseek": ("DEEPSEEK_API_BASE",), + "docker_model_runner": ("DOCKER_MODEL_RUNNER_API_BASE",), + "empower": ("EMPOWER_API_BASE",), + "featherless_ai": ("FEATHERLESS_AI_API_BASE", "FEATHERLESS_API_BASE"), + "fireworks_ai": ("FIREWORKS_API_BASE",), + "friendliai": ("FRIENDLI_API_BASE",), + "galadriel": ("GALADRIEL_API_BASE",), + "gdc": ("GDC_API_BASE",), + "gemini": ("GEMINI_API_BASE",), + "github": ("GITHUB_API_BASE",), + "github_copilot": ("GITHUB_COPILOT_API_BASE",), + "gigachat": ("GIGACHAT_API_BASE",), + "gradient_ai": ("GRADIENT_AI_AGENT_ENDPOINT",), + "groq": ("GROQ_API_BASE",), + "heroku": ("HEROKU_API_BASE",), + "hosted_vllm": ("HOSTED_VLLM_API_BASE",), + "huggingface": ("HF_API_BASE", "HUGGINGFACE_API_BASE"), + "hyperbolic": ("HYPERBOLIC_API_BASE",), + "inception": ("INCEPTION_API_BASE",), + "lambda_ai": ("LAMBDA_API_BASE",), + "langflow": ("LANGFLOW_API_BASE",), + "langgraph": ("LANGGRAPH_API_BASE",), + "lemonade": ("LEMONADE_API_BASE",), + "litellm_proxy": ("LITELLM_PROXY_API_BASE",), + "llamafile": ("LLAMAFILE_API_BASE",), + "lm_studio": ("LM_STUDIO_API_BASE",), + "manus": ("MANUS_API_BASE",), + "maritalk": ("MARITALK_API_BASE",), + "meta_llama": ("LLAMA_API_BASE",), + "minimax": ("MINIMAX_API_BASE",), + "mistral": ("MISTRAL_AZURE_API_BASE",), + "modelscope": ("MODELSCOPE_API_BASE",), + "moonshot": ("MOONSHOT_API_BASE",), + "morph": ("MORPH_API_BASE",), + "nebius": ("NEBIUS_API_BASE",), + "novita": ("NOVITA_API_BASE",), + "nscale": ("NSCALE_API_BASE",), + "nvidia_nim": ("NVIDIA_NIM_API_BASE",), + "nvidia_riva": ("NVIDIA_RIVA_API_BASE",), + "nlp_cloud": ("NLP_CLOUD_API_BASE",), + "ollama": ("OLLAMA_API_BASE",), + "openai_like": ("OPENAI_LIKE_API_BASE",), + "openrouter": ("OPENROUTER_API_BASE",), + "ovhcloud": ("OVHCLOUD_API_BASE",), + "perplexity": ("PERPLEXITY_API_BASE",), + "predibase": ("PREDIBASE_API_BASE",), + "ragflow": ("RAGFLOW_API_BASE",), + "replicate": ("REPLICATE_API_BASE",), + "sambanova": ("SAMBANOVA_API_BASE",), + "tencent": ("TENCENT_API_BASE",), + "together_ai": ("TOGETHER_AI_API_BASE",), + "v0": ("V0_API_BASE",), + "vercel_ai_gateway": ("VERCEL_AI_GATEWAY_API_BASE",), + "vertex_ai": ("VERTEXAI_API_BASE",), + "volcengine": ("VOLCENGINE_API_BASE",), + "wandb": ("WANDB_API_BASE",), + "xai": ("XAI_API_BASE",), + "xinference": ("XINFERENCE_API_BASE",), + "zai": ("ZAI_API_BASE",), +} + +_WATSONX_ROUTING_ENV_VARS = { + "api_base": ("WATSONX_API_BASE", "WATSONX_URL", "WX_URL", "WML_URL"), + "project_id": ("WATSONX_PROJECT_ID", "WX_PROJECT_ID", "PROJECT_ID"), + "space_id": ("WATSONX_DEPLOYMENT_SPACE_ID", "WATSONX_SPACE_ID", "WX_SPACE_ID", "SPACE_ID"), + "region_name": ("WATSONX_REGION", "WX_REGION", "REGION"), + "token": ("WATSONX_TOKEN",), + "zen_api_key": ("WATSONX_ZENAPIKEY",), +} + +PROVIDER_ROUTING_ENV_VARS = { + "azure": { + "api_base": ("AZURE_API_BASE", "AZURE_OPENAI_ENDPOINT"), + "api_version": ("AZURE_API_VERSION",), + }, + "bedrock": { + "aws_bedrock_runtime_endpoint": ("AWS_BEDROCK_RUNTIME_ENDPOINT",), + "aws_region_name": ("AWS_REGION_NAME", "AWS_REGION", "AWS_DEFAULT_REGION"), + }, + "bedrock_mantle": { + "aws_bedrock_runtime_endpoint": ("AWS_BEDROCK_RUNTIME_ENDPOINT",), + "aws_region_name": ("BEDROCK_MANTLE_REGION", "AWS_REGION_NAME", "AWS_REGION", "AWS_DEFAULT_REGION"), + }, + "cloudflare": { + "api_base": ("CLOUDFLARE_API_BASE", "CLOUDFLARE_ACCOUNT_ID"), + }, + "openai": { + "organization": ("OPENAI_ORGANIZATION",), + }, + "openrouter": {"api_base": ("OPENROUTER_API_BASE",)}, + "vertex_ai": { + "vertex_project": ("VERTEXAI_PROJECT", "GOOGLE_CLOUD_PROJECT", "GCLOUD_PROJECT"), + "vertex_location": ("VERTEXAI_LOCATION", "VERTEX_LOCATION"), + }, + "watsonx": dict(_WATSONX_ROUTING_ENV_VARS), + "watsonx_text": dict(_WATSONX_ROUTING_ENV_VARS), +} + +# Keep aliases in the same precedence order as LiteLLM 1.100.0. +PROVIDER_API_KEY_ENV_VARS = { + "ai21": ("AI21_API_KEY",), + "ai21_chat": ("AI21_API_KEY",), + "aiml": ("AIML_API_KEY",), + "aleph_alpha": ("ALEPH_ALPHA_API_KEY", "ALEPHALPHA_API_KEY"), + "amazon_nova": ("AMAZON_NOVA_API_KEY",), + "anyscale": ("ANYSCALE_API_KEY",), + "anthropic": ("ANTHROPIC_API_KEY",), + "azure": ("AZURE_OPENAI_API_KEY", "AZURE_API_KEY"), + "azure_ai": ("AZURE_AI_API_KEY",), + "baseten": ("BASETEN_API_KEY",), + "bedrock": ("AWS_BEARER_TOKEN_BEDROCK",), + "bedrock_mantle": ("BEDROCK_MANTLE_API_KEY", "AWS_BEARER_TOKEN_BEDROCK"), + "bytez": ("BYTEZ_API_KEY",), + "cerebras": ("CEREBRAS_API_KEY",), + "clarifai": ("CLARIFAI_API_KEY",), + "cloudflare": ("CLOUDFLARE_API_KEY",), + "codestral": ("CODESTRAL_API_KEY",), + "cohere": ("COHERE_API_KEY", "CO_API_KEY"), + "cohere_chat": ("COHERE_API_KEY", "CO_API_KEY"), + "cometapi": ("COMETAPI_KEY",), + "compactifai": ("COMPACTIFAI_API_KEY",), + "custom_openai": ("OPENAI_API_KEY",), # Preserve the existing OpenAI-compatible fallback. + "dashscope": ("DASHSCOPE_API_KEY",), + "databricks": ("DATABRICKS_API_KEY",), + "datarobot": ("DATAROBOT_API_TOKEN",), + "deepinfra": ("DEEPINFRA_API_KEY",), + "deepseek": ("DEEPSEEK_API_KEY",), + "docker_model_runner": ("DOCKER_MODEL_RUNNER_API_KEY",), + "empower": ("EMPOWER_API_KEY",), + "featherless_ai": ("FEATHERLESS_AI_API_KEY", "FEATHERLESS_API_KEY"), + "fireworks_ai": ( + "FIREWORKS_API_KEY", + "FIREWORKS_AI_API_KEY", + "FIREWORKSAI_API_KEY", + "FIREWORKS_AI_TOKEN", + ), + "friendliai": ("FRIENDLIAI_API_KEY", "FRIENDLI_TOKEN"), + "galadriel": ("GALADRIEL_API_KEY",), + "gdc": ("GDC_API_KEY",), + "gemini": ("GOOGLE_API_KEY", "GEMINI_API_KEY", "PALM_API_KEY"), + "github": ("GITHUB_API_KEY",), + "gigachat": ("GIGACHAT_API_KEY", "GIGACHAT_CREDENTIALS"), + "gradient_ai": ("GRADIENT_AI_API_KEY",), + "groq": ("GROQ_API_KEY",), + "heroku": ("HEROKU_API_KEY",), + "hosted_vllm": ("HOSTED_VLLM_API_KEY",), + "huggingface": ("HF_TOKEN", "HUGGINGFACE_API_KEY"), + "hyperbolic": ("HYPERBOLIC_API_KEY",), + "inception": ("INCEPTION_API_KEY",), + "lambda_ai": ("LAMBDA_API_KEY",), + "langflow": ("LANGFLOW_API_KEY",), + "langgraph": ("LANGGRAPH_API_KEY",), + "lemonade": ("LEMONADE_API_KEY",), + "litellm_proxy": ("LITELLM_PROXY_API_KEY",), + "llamafile": ("LLAMAFILE_API_KEY",), + "lm_studio": ("LM_STUDIO_API_KEY",), + "manus": ("MANUS_API_KEY",), + "maritalk": ("MARITALK_API_KEY",), + "meta_llama": ("LLAMA_API_KEY",), + "minimax": ("MINIMAX_API_KEY",), + "mistral": ("MISTRAL_AZURE_API_KEY", "MISTRAL_API_KEY"), + "modelscope": ("MODELSCOPE_API_KEY",), + "moonshot": ("MOONSHOT_API_KEY",), + "morph": ("MORPH_API_KEY",), + "nebius": ("NEBIUS_API_KEY",), + "novita": ("NOVITA_API_KEY",), + "nscale": ("NSCALE_API_KEY",), + "nvidia_nim": ("NVIDIA_NIM_API_KEY",), + "nlp_cloud": ("NLP_CLOUD_API_KEY",), + "ollama": ("OLLAMA_API_KEY",), + "openai": ("OPENAI_API_KEY",), + "openai_like": ("OPENAI_LIKE_API_KEY",), + "openrouter": ("OPENROUTER_API_KEY", "OR_API_KEY"), + "ovhcloud": ("OVHCLOUD_API_KEY",), + "perplexity": ("PERPLEXITYAI_API_KEY", "PERPLEXITY_API_KEY"), + "predibase": ("PREDIBASE_API_KEY",), + "ragflow": ("RAGFLOW_API_KEY",), + "replicate": ("REPLICATE_API_KEY", "REPLICATE_API_TOKEN"), + "sambanova": ("SAMBANOVA_API_KEY",), + "sap": ("AICORE_SERVICE_KEY",), + "snowflake": ("SNOWFLAKE_JWT",), + "tencent": ("TENCENT_API_KEY",), + "text-completion-codestral": ("CODESTRAL_API_KEY",), + "together_ai": ( + "TOGETHER_API_KEY", + "TOGETHER_AI_API_KEY", + "TOGETHERAI_API_KEY", + "TOGETHER_AI_TOKEN", + ), + "v0": ("V0_API_KEY",), + "vercel_ai_gateway": ("VERCEL_AI_GATEWAY_API_KEY", "VERCEL_OIDC_TOKEN"), + "volcengine": ("VOLCENGINE_API_KEY",), + "wandb": ("WANDB_API_KEY",), + "watsonx": ("WATSONX_APIKEY", "WATSONX_API_KEY", "WX_API_KEY", "WATSONX_ZENAPIKEY"), + "watsonx_text": ("WATSONX_APIKEY", "WATSONX_API_KEY", "WX_API_KEY"), + "xai": ("XAI_API_KEY",), + "xiaomi_mimo": ("XIAOMI_MIMO_API_KEY",), + "xinference": ("XINFERENCE_API_KEY",), + "zai": ("ZAI_API_KEY",), +} + +OPENAI_COMPATIBLE_REQUEST_PROVIDERS = {"custom_openai", "openai_like"} +MANAGED_AUTH_REQUEST_PROVIDERS = {"chatgpt", "github_copilot"} +OPENAI_RAW_HTTP_REQUEST_PROVIDERS = { + "aiohttp_openai", + "azure_ai", + "cometapi", + "deepseek", + "fireworks_ai", + "groq", + "heroku", + "hosted_vllm", + "minimax", + "openai_like", + "openrouter", + "ragflow", + "together_ai", + "xai", +} +AZURE_AD_TOKEN_ENV_VARS = ("AZURE_AD_TOKEN", "AZURE_OPENAI_AD_TOKEN") +AZURE_OIDC_ENV_VARS = ("AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_AUTHORITY_HOST", "AZURE_SCOPE") +AZURE_OIDC_AUTH_ENV_VARS = ("AZURE_CLIENT_SECRET", "AZURE_USERNAME", "AZURE_PASSWORD") + +PROVIDER_API_KEY_GLOBALS = { + "ai21": ("ai21_key",), + "ai21_chat": ("ai21_key",), + "aleph_alpha": ("aleph_alpha_key",), + "amazon_nova": ("amazon_nova_api_key",), + "anthropic": ("anthropic_key",), + "azure": ("azure_key",), + "baseten": ("baseten_key",), + "bytez": ("bytez_key",), + "cloudflare": ("cloudflare_api_key",), + "cohere": ("cohere_key",), + "cohere_chat": ("cohere_key",), + "cometapi": ("cometapi_key",), + "custom_openai": ("openai_key",), + "databricks": ("databricks_key",), + "gdc": ("gdc_key",), + "gigachat": ("gigachat_key",), + "groq": ("groq_key",), + "huggingface": ("huggingface_key",), + "inception": ("inception_key",), + "lemonade": ("lemonade_key",), + "maritalk": ("maritalk_key",), + "nebius": ("nebius_key",), + "nlp_cloud": ("nlp_cloud_key",), + "ollama": ("ollama_key", "openai_key"), + "openai": ("openai_key",), + "openai_like": ("openai_like_key",), + "openrouter": ("openrouter_key",), + "ovhcloud": ("ovhcloud_key",), + "predibase": ("predibase_key",), + "replicate": ("replicate_key",), + "sap": ("sap_service_key",), + "together_ai": ("togetherai_api_key",), + "wandb": ("wandb_key",), + "xai": ("xai_key",), +} + + +def _require_litellm_interface(interface, name: str, methods=()): + """Reject unavailable isolation interfaces without disabling their protection.""" + available = ( + all(callable(getattr(interface, method, None)) for method in methods) + if methods else callable(interface) + ) + if not available: + raise RuntimeError( + f"LiteLLM interface {name} is unavailable; request isolation cannot be guaranteed. " + "Restore the project's locked LiteLLM version or update and validate this integration." + ) + + +_azure_oidc_request = ContextVar("pr_agent_azure_oidc_request", default=None) +_azure_ad_responses_request = ContextVar("pr_agent_azure_ad_responses_request", default=None) +_azure_oidc_bridge = None + + +def _exchange_azure_oidc_token(selector, environment, native=None): + """Bind the native outer exchange to captured inputs without mutating its globals.""" + from litellm.llms.azure.common_utils import get_azure_ad_token_from_oidc + + if native is None: + native = get_azure_ad_token_from_oidc + if _azure_oidc_bridge is not None and native is _azure_oidc_bridge["exchange"]: + native = _azure_oidc_bridge["original_exchange"] + if set(environment) != set(AZURE_OIDC_ENV_VARS): + raise RuntimeError("Azure OIDC exchange snapshot is incomplete") + if not isinstance(native, FunctionType): + raise RuntimeError("LiteLLM's Azure OIDC exchange interface is incompatible") + parameters = tuple(inspect.signature(native).parameters.values()) + names = ("azure_ad_token", "azure_client_id", "azure_tenant_id", "scope") + if len(parameters) != len(names) or any( + parameter.name != name + or parameter.kind != inspect.Parameter.POSITIONAL_OR_KEYWORD + or parameter.default is not (inspect.Parameter.empty if index == 0 else None) + for index, (parameter, name) in enumerate(zip(parameters, names, strict=True)) + ): + raise RuntimeError("LiteLLM's Azure OIDC exchange interface is incompatible") + native_globals = dict(native.__globals__) + cache = native_globals.get("azure_ad_cache") + _require_litellm_interface(cache, "Azure OIDC cache", ("get_cache", "set_cache")) + _require_litellm_interface(native_globals.get("get_secret_str"), "Azure OIDC assertion resolver") + + def captured_getenv(name, default=None): + if name not in environment: + raise RuntimeError("LiteLLM's Azure OIDC exchange uses an uncaptured input") + value = environment[name] + return default if value is None else value + + # Match the pinned helper's default without normalizing a deliberate empty scope. + scope = environment["AZURE_SCOPE"] + effective_scope = "https://cognitiveservices.azure.com/.default" if scope is None else scope + + def scoped_key(key): + return json.dumps(["pr-agent.azure-oidc.v1", effective_scope, key]) + + def get_cache(key, *args, **kwargs): + return cache.get_cache(scoped_key(key), *args, **kwargs) + + def set_cache(key, *args, **kwargs): + return cache.set_cache(scoped_key(key), *args, **kwargs) + + # Keep native assertion resolution (including deployment-owned source selection) + # and refresh/TTL behavior. Only the outer exchange's inputs are request-owned. + native_globals["os"] = SimpleNamespace(getenv=captured_getenv) + native_globals["azure_ad_cache"] = SimpleNamespace(get_cache=get_cache, set_cache=set_cache) + exchange = FunctionType(native.__code__, native_globals, native.__name__, native.__defaults__, native.__closure__) + exchange.__kwdefaults__ = copy.copy(native.__kwdefaults__) + return exchange( + selector, + azure_client_id=environment["AZURE_CLIENT_ID"], + azure_tenant_id=environment["AZURE_TENANT_ID"], + scope=scope, + ) + + +@lru_cache(maxsize=128) +def _azure_oidc_entra_provider(tenant_id, client_id, client_secret, scope, authority): + from azure.identity import ClientSecretCredential, get_bearer_token_provider + + credential = ClientSecretCredential(tenant_id, client_id, client_secret, authority=authority) + return get_bearer_token_provider(credential, scope) + + +def _azure_oidc_companion_globals(native_globals, context): + """Bind companion factories only inside the selected native auth path.""" + native_globals = dict(native_globals) + original_entra = native_globals["get_azure_ad_token_from_entra_id"] + original_password = native_globals["get_azure_ad_token_from_username_password"] + + def authority(): + from azure.identity import AzureAuthorityHosts + + value = context["environment"]["AZURE_AUTHORITY_HOST"] + if value == "": + raise ValueError("Azure companion credentials require a nonempty captured authority") + return AzureAuthorityHosts.AZURE_PUBLIC_CLOUD if value is None else value + + def entra(tenant_id, client_id, client_secret, scope="https://cognitiveservices.azure.com/.default"): + if not isinstance(original_entra, FunctionType): + raise RuntimeError("LiteLLM's Azure Entra factory is incompatible") + captured_authority = authority() + + def cached(tenant_id, client_id, client_secret, scope): + return _azure_oidc_entra_provider(tenant_id, client_id, client_secret, scope, captured_authority) + + # Retain native secret-selector resolution, but never use the native + # authority-free provider cache. Credentials retain authority on refresh. + globals_view = {**original_entra.__globals__, "_cached_entra_id_token_provider": cached} + factory = FunctionType( + original_entra.__code__, globals_view, original_entra.__name__, + original_entra.__defaults__, original_entra.__closure__, + ) + factory.__kwdefaults__ = copy.copy(original_entra.__kwdefaults__) + return factory(tenant_id, client_id, client_secret, scope) + + def password(client_id, azure_username, azure_password, scope="https://cognitiveservices.azure.com/.default"): + from azure.identity import UsernamePasswordCredential, get_bearer_token_provider + + if not isinstance(original_password, FunctionType): + raise RuntimeError("LiteLLM's Azure password factory is incompatible") + # The native password path leaves tenant selection at the SDK default. + credential = UsernamePasswordCredential( + client_id=client_id, username=azure_username, password=azure_password, authority=authority(), + ) + return get_bearer_token_provider(credential, scope) + + native_globals["get_azure_ad_token_from_entra_id"] = entra + native_globals["get_azure_ad_token_from_username_password"] = password + + def reject_implicit_refresh(*args, **kwargs): + # Native discovery re-reads ambient identity after companion selection. + # ValueError is swallowed by LiteLLM and would silently change auth. + raise RuntimeError("Request isolation forbids implicit Azure AD credential discovery; " + "configure complete companion credentials instead") + + native_globals["get_azure_ad_token_provider"] = reject_implicit_refresh + return native_globals + + +def _install_azure_oidc_bridge(): + """Bind native OIDC and SDK AD selection without changing key precedence.""" + global _azure_oidc_bridge + from litellm.llms.azure import azure as azure_module + from litellm.llms.azure import common_utils as azure_common + from litellm.llms.openai.common_utils import BaseOpenAILLM + + if _azure_oidc_bridge is not None: + if ( + azure_common.get_azure_ad_token_from_oidc is not _azure_oidc_bridge["exchange"] + or azure_module.get_azure_ad_token_from_oidc is not _azure_oidc_bridge["exchange"] + or azure_common.get_azure_ad_token is not _azure_oidc_bridge["resolver"] + or BaseOpenAILLM.get_openai_client_cache_key is not _azure_oidc_bridge["cache_key"] + or azure_common.BaseAzureLLM._resolve_env_var is not _azure_oidc_bridge["resolve_env"] + or azure_common.BaseAzureLLM.get_azure_openai_client is not _azure_oidc_bridge["client"] + or azure_common.BaseAzureLLM._base_validate_azure_environment is not _azure_oidc_bridge["validate"] + or azure_common.BaseAzureLLM.initialize_azure_sdk_client is not _azure_oidc_bridge["initialize"] + ): + raise RuntimeError("LiteLLM's Azure OIDC bridge interfaces were replaced") + return + original_exchange = azure_common.get_azure_ad_token_from_oidc + original_resolver = azure_common.get_azure_ad_token + original_cache_key = BaseOpenAILLM.get_openai_client_cache_key + original_resolve_env = azure_common.BaseAzureLLM._resolve_env_var + original_client = azure_common.BaseAzureLLM.get_azure_openai_client + original_validate = azure_common.BaseAzureLLM._base_validate_azure_environment + original_initialize = azure_common.BaseAzureLLM.initialize_azure_sdk_client + client_signature = inspect.signature(original_client) + initialize_signature = inspect.signature(original_initialize) + if ( + not all(isinstance(function, FunctionType) for function in ( + original_exchange, original_resolver, original_cache_key, original_resolve_env, + original_client, original_validate, original_initialize, + )) + or azure_module.get_azure_ad_token_from_oidc is not original_exchange + ): + raise RuntimeError("LiteLLM's Azure OIDC interfaces are incompatible") + + @wraps(original_exchange) + def exchange(azure_ad_token, azure_client_id=None, azure_tenant_id=None, scope=None): + context = _azure_oidc_request.get() + if context is not None and context["selector"] and azure_ad_token == context["selector"]: + return _exchange_azure_oidc_token(azure_ad_token, context["environment"], original_exchange) + return original_exchange(azure_ad_token, azure_client_id, azure_tenant_id, scope) + + def resolve_captured(litellm_params, context): + native_globals = _azure_oidc_companion_globals(original_resolver.__globals__, context) + native_os = native_globals["os"] + native_get_secret = native_globals["get_secret_str"] + + def get_secret_str(name, *args, **kwargs): + if name == "AZURE_AD_TOKEN": + return context["dispatch_token"] + return native_get_secret(name, *args, **kwargs) + + def getenv(name, default=None): + if name in context["auth_environment"]: + value = context["auth_environment"][name] + return default if value is None else value + return native_os.getenv(name, default) + + native_globals["os"] = SimpleNamespace(getenv=getenv) + native_globals["get_secret_str"] = get_secret_str + native_globals["get_azure_ad_token_from_oidc"] = exchange + bound = FunctionType( + original_resolver.__code__, native_globals, original_resolver.__name__, + original_resolver.__defaults__, original_resolver.__closure__, + ) + bound.__kwdefaults__ = copy.copy(original_resolver.__kwdefaults__) + return bound(litellm_params) + + @wraps(original_resolver) + def resolver(litellm_params): + ad_context = _azure_ad_responses_request.get() + if ad_context is not None and litellm_params.get("azure_ad_token") == ad_context["dispatch_token"]: + return resolve_captured(litellm_params, ad_context) + context = _azure_oidc_request.get() + if context is None or not context["selector"] or litellm_params.get("azure_ad_token") != context["selector"]: + return original_resolver(litellm_params) + environment = context["environment"] + if not environment["AZURE_CLIENT_ID"] or not environment["AZURE_TENANT_ID"]: + raise ValueError("Azure OIDC requires captured client and tenant identifiers") + return resolve_captured(litellm_params, context) + + @wraps(original_resolve_env) + def resolve_env(owner, litellm_params, param_key, env_var_key): + context = _azure_oidc_request.get() + if ( + context is not None and litellm_params.get("azure_ad_token") == context["dispatch_token"] + and env_var_key in context["auth_environment"] + ): + value = litellm_params.get(param_key) + return value if value is not None else context["auth_environment"][env_var_key] + return original_resolve_env(owner, litellm_params, param_key, env_var_key) + + @wraps(original_initialize) + def initialize(*args, **kwargs): + context = _azure_oidc_request.get() + if context is None: + return original_initialize(*args, **kwargs) + bound = initialize_signature.bind(*args, **kwargs) + params = bound.arguments.get("litellm_params") or {} + if params.get("azure_ad_token") != context["dispatch_token"]: + return original_initialize(*args, **kwargs) + native_globals = _azure_oidc_companion_globals(original_initialize.__globals__, context) + factory = FunctionType( + original_initialize.__code__, native_globals, original_initialize.__name__, + original_initialize.__defaults__, original_initialize.__closure__, + ) + factory.__kwdefaults__ = copy.copy(original_initialize.__kwdefaults__) + return factory(*args, **kwargs) + + @wraps(original_cache_key) + def cache_key(client_initialization_params, client_type): + key = original_cache_key(client_initialization_params, client_type) + context = _azure_oidc_request.get() + if ( + context is not None and client_type == "azure" + and client_initialization_params.get("azure_ad_token") == context["selector_hash"] + ): + return f"{key}|pr_agent_oidc={context['identity_hash']}" + return key + + @wraps(original_client) + def client(*args, **kwargs): + context = _azure_oidc_request.get() + if context is None or not context["generated_guard"]: + return original_client(*args, **kwargs) + bound = client_signature.bind(*args, **kwargs) + params = bound.arguments.get("litellm_params") or {} + if ( + bound.arguments.get("api_key") == DUMMY_LITELLM_API_KEY + and params.get("azure_ad_token") == context["dispatch_token"] + ): + # Dispatch has already resolved fallbacks. Restore native keyless + # selection before the SDK initializer and its client-cache lookup. + bound.arguments["api_key"] = None + return original_client(*bound.args, **bound.kwargs) + + @wraps(original_validate) + def validate(headers, litellm_params): + ad_context = _azure_ad_responses_request.get() + if ( + ad_context is not None and ad_context["generated_guard"] and "api-key" not in headers + and litellm_params.get("api_key") == DUMMY_LITELLM_API_KEY + and litellm_params.get("azure_ad_token") == ad_context["dispatch_token"] + ): + # Only a positively identified generated guard may bypass key + # selection. Retain native AD provider precedence, not global keys. + headers = dict(headers) + token = resolve_captured(litellm_params, ad_context) + if token: + headers["Authorization"] = f"Bearer {token}" + return headers + context = _azure_oidc_request.get() + if ( + context is None or not context["selector"] or not context["generated_guard"] or "api-key" in headers + or litellm_params.get("api_key") != DUMMY_LITELLM_API_KEY + or litellm_params.get("azure_ad_token") != context["dispatch_token"] + ): + return original_validate(headers, litellm_params) + # Do not retry process-wide API-key fallback with a cleared guard. + headers = dict(headers) + token = resolver(litellm_params) + if token: + headers["Authorization"] = f"Bearer {token}" + return headers + + azure_common.get_azure_ad_token_from_oidc = exchange + azure_module.get_azure_ad_token_from_oidc = exchange + azure_common.get_azure_ad_token = resolver + azure_common.BaseAzureLLM._resolve_env_var = resolve_env + azure_common.BaseAzureLLM.get_azure_openai_client = client + azure_common.BaseAzureLLM._base_validate_azure_environment = staticmethod(validate) + azure_common.BaseAzureLLM.initialize_azure_sdk_client = initialize + BaseOpenAILLM.get_openai_client_cache_key = staticmethod(cache_key) + _azure_oidc_bridge = { + "exchange": exchange, "resolver": resolver, "cache_key": cache_key, + "original_exchange": original_exchange, + "resolve_env": resolve_env, + "client": client, "validate": validate, "initialize": initialize, + } + + +def _is_openai_compatible_request_provider(provider: str) -> bool: + return provider in getattr(litellm, "openai_compatible_providers", ()) or JSONProviderRegistry.exists(provider) + + +def _uses_openai_text_completion_transport(model: str | None, provider: str | None) -> bool: + """Return whether LiteLLM will dispatch this request through the OpenAI text-completion SDK.""" + if provider == "text-completion-openai": + return True + if not isinstance(model, str): + return False + return any(prefix in model for prefix in ("ft:babbage-002", "ft:davinci-002")) + + +def _uses_openai_responses_transport(model: str | None, provider: str | None) -> bool: + """Return whether LiteLLM will send this model through its raw Responses transport.""" + if not isinstance(model, str) or _uses_openai_text_completion_transport(model, provider): + return False + provider = provider or "" + canonical_provider = PROVIDER_SETTING_ALIASES.get(provider, provider) + if "/" in model and model.split("/", 1)[0] in (provider, canonical_provider): + model = model.split("/", 1)[1] + if model.startswith("responses/"): + return True + if provider == "openai" and getattr(litellm, "route_all_chat_openai_to_responses", False): + return True + _require_litellm_interface(_get_model_info_helper, "_get_model_info_helper") + try: + model_info = _get_model_info_helper(model=model, custom_llm_provider=provider) + except Exception: + return False + return model_info.get("mode") == "responses" + + +def _uses_provider_api_key(provider: str) -> bool: + """Return whether PR-Agent should forward or guard this provider's native API key.""" + return provider in PROVIDER_API_KEY_ENV_VARS or JSONProviderRegistry.exists(provider) + + +_SDK_HEADER_MARKER = "x-pr-agent-sdk-header-snapshot" +_sdk_request_headers = ContextVar("pr_agent_sdk_request_headers", default=None) + + +class _CapturedSDKHeader(openai.Omit): + """Distinguish generated account guards from explicit header omissions.""" + + +class _SDKHeaderSnapshot: + # Keep credential-bearing snapshots out of SDK DEBUG and fallback serialization. + __slots__ = ("snapshot",) + + def __init__(self, snapshot): + self.snapshot = snapshot + + def __repr__(self): + return "" + + +def _merge_sdk_headers(*layers): + """Merge before HTTPX so differently cased names cannot duplicate credentials.""" + result = {} + names = {} + for layer in layers: + for name, value in layer.items(): + lower = name.lower() + canonical = names.setdefault(lower, name) + result[canonical] = value + return result + + +def _check_sdk_marker_collision(headers): + if any(name.lower() == _SDK_HEADER_MARKER for name in headers): + raise ValueError("Request headers contain a reserved SDK isolation header") + + +class _SDKHeaderLoggingProxy: + def __init__(self, logger): + object.__setattr__(self, "_logger", logger) + + def __getattr__(self, name): + return getattr(self._logger, name) + + def __setattr__(self, name, value): + setattr(self._logger, name, value) + + def pre_call(self, *args, **kwargs): + additional = kwargs.get("additional_args") + if isinstance(additional, dict) and isinstance(additional.get("complete_input_dict"), dict): + data = dict(additional["complete_input_dict"]) + if isinstance(data.get("extra_headers"), dict): + data["extra_headers"] = { + key: value for key, value in data["extra_headers"].items() + if key.lower() != _SDK_HEADER_MARKER + } + kwargs["additional_args"] = {**additional, "complete_input_dict": data} + return self._logger.pre_call(*args, **kwargs) + + +def _sdk_data_adapter(original, *, text=False, copilot=False): + signature = inspect.signature(original) + + @wraps(original) + def scoped(*args, **kwargs): + snapshot = _sdk_request_headers.get() + if snapshot is None: + return original(*args, **kwargs) + bound = signature.bind(*args, **kwargs) + if copilot: + if snapshot["provider"] != "github_copilot": + return original(*args, **kwargs) + params = bound.arguments.get("optional_params") or {} + bound.arguments["headers"] = _merge_sdk_headers( + params.get("extra_headers") or {}, bound.arguments.get("headers") or {}, + ) + else: + data = dict(bound.arguments["data"]) + if text and snapshot["provider"] == "azure": + # LiteLLM also copies this auth argument into text request data. + # Keep it on the native client parameters, never in the SDK payload. + data.pop("azure_ad_token", None) + headers = dict(data.get("extra_headers") or {}) + _check_sdk_marker_collision(headers) + headers[_SDK_HEADER_MARKER] = _SDKHeaderSnapshot(snapshot) + data["extra_headers"] = headers + bound.arguments["data"] = data + if text and bound.arguments.get("logging_obj") is not None: + bound.arguments["logging_obj"] = _SDKHeaderLoggingProxy(bound.arguments["logging_obj"]) + # Capture before returning a lazy text streaming generator, not at first iteration. + return original(*bound.args, **bound.kwargs) + + scoped._pr_agent_sdk_headers_original = original + scoped._pr_agent_sdk_headers_context = _sdk_request_headers + return scoped + + +def _install_sdk_header_bridge(): + """Adapt the locked SDK's request boundary without mutating shared clients.""" + try: + from litellm.llms.azure.azure import AzureChatCompletion + from litellm.llms.azure.completion.handler import AzureTextCompletion + from litellm.llms.openai.completion.handler import OpenAITextCompletion + from litellm.llms.openai.openai import OpenAIChatCompletion + from openai._base_client import BaseClient + except ImportError as error: + raise RuntimeError("The installed LiteLLM/OpenAI SDK cannot isolate request headers") from error + + bindings = ( + (OpenAIChatCompletion, "make_openai_chat_completion_request", False, False), + (AzureChatCompletion, "make_azure_openai_chat_completion_request", False, False), + (OpenAITextCompletion, "acompletion", True, False), + (OpenAITextCompletion, "async_streaming", True, False), + (AzureTextCompletion, "acompletion", True, False), + (AzureTextCompletion, "async_streaming", True, False), + (OpenAIChatCompletion, "completion", False, True), + ) + _require_litellm_interface(BaseClient, "OpenAI.BaseClient", ("_build_headers",)) + for owner, name, _, copilot in bindings: + _require_litellm_interface(owner, owner.__name__, (name,)) + required = {"optional_params", "headers"} if copilot else {"data", "logging_obj"} + if not required <= inspect.signature(getattr(owner, name)).parameters.keys(): + raise RuntimeError("The installed LiteLLM SDK request interface cannot isolate headers") + + current = BaseClient._build_headers + if getattr(current, "_pr_agent_sdk_headers_context", None) is not _sdk_request_headers: + sdk_build_headers = getattr(current, "_pr_agent_sdk_headers_original", current) + + @wraps(sdk_build_headers) + def build_headers(client, options, *, retries_taken=0): + headers = dict(options.headers or {}) + marker = headers.get(_SDK_HEADER_MARKER) + if not isinstance(marker, _SDKHeaderSnapshot): + return sdk_build_headers(client, options, retries_taken=retries_taken) + if type(client) not in (openai.OpenAI, openai.AsyncOpenAI, openai.AzureOpenAI, openai.AsyncAzureOpenAI): + raise ValueError("Custom SDK clients are incompatible with request header isolation") + del headers[_SDK_HEADER_MARKER] + snapshot = marker.snapshot + view = copy.copy(client) + view._custom_headers = {} + view.organization = snapshot["organization"] + view.project = snapshot["project"] + explicit = dict(snapshot["explicit_headers"]) + for layer in (headers, explicit): + for name, value in tuple(layer.items()): + if isinstance(value, _CapturedSDKHeader): + # Captured SDK defaults already provide the account identity; + # internal guards must not override native custom headers. + del layer[name] + native = {**view._auth_headers(options.security), **view.default_headers} + merged = _merge_sdk_headers(native, snapshot["custom_headers"], headers, explicit) + # Native defaults retain their canonical spelling; options override them without duplicates. + request_options = options.model_copy(update={"headers": merged}) + result = sdk_build_headers(view, request_options, retries_taken=retries_taken) + if _SDK_HEADER_MARKER in result: + raise RuntimeError("SDK isolation marker reached HTTP headers") + return result + + build_headers._pr_agent_sdk_headers_original = sdk_build_headers + build_headers._pr_agent_sdk_headers_context = _sdk_request_headers + BaseClient._build_headers = build_headers + for owner, name, text, copilot in bindings: + current = getattr(owner, name) + if getattr(current, "_pr_agent_sdk_headers_context", None) is not _sdk_request_headers: + original = getattr(current, "_pr_agent_sdk_headers_original", current) + setattr(owner, name, _sdk_data_adapter(original, text=text, copilot=copilot)) + + +def _azure_ai_native_transport(provider: str, model: str | None) -> tuple[str, str | None]: + """Follow Azure AI's native model classification without resolving credentials.""" + if provider == "azure_ai" and isinstance(model, str): + from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig + + native_model = model.removeprefix("azure_ai/") + if AzureAIStudioConfig()._is_azure_openai_model(native_model, None): + return "azure", native_model + return provider, model + + +def _is_cloudflare_gateway(api_base: str | None) -> bool: + """Match the AI Gateway host itself; a substring test also matches it inside a path or query.""" + hostname = urlparse(api_base or "").hostname or "" + return hostname == "gateway.ai.cloudflare.com" or hostname.endswith(".gateway.ai.cloudflare.com") + + +def _request_local_openai_headers(provider: str, organization=None, model: str | None = None) -> dict | None: + """Block LiteLLM's process-wide OpenAI headers for one request.""" + transport_provider, model = _azure_ai_native_transport(provider, model) + provider = PROVIDER_SETTING_ALIASES.get(transport_provider, transport_provider) + if not ( + provider in ("azure", "openai", "openrouter") + or provider in OPENAI_COMPATIBLE_REQUEST_PROVIDERS + or provider in MANAGED_AUTH_REQUEST_PROVIDERS + or _is_openai_compatible_request_provider(provider) + ): + return None + if _uses_openai_responses_transport(model, transport_provider): + if provider == "openai" and organization: + return {"OpenAI-Organization": organization} + return None + experimental_raw_http_handler = os.environ.get( + "EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "" + ).strip().lower() == "true" + uses_raw_http_handler = transport_provider in OPENAI_RAW_HTTP_REQUEST_PROVIDERS or ( + experimental_raw_http_handler + and not _uses_openai_text_completion_transport(model, transport_provider) + and ( + transport_provider == "openai" + or transport_provider in OPENAI_COMPATIBLE_REQUEST_PROVIDERS + or transport_provider in MANAGED_AUTH_REQUEST_PROVIDERS + or _is_openai_compatible_request_provider(transport_provider) + ) + ) + if uses_raw_http_handler: + return None + request_organization = organization if provider == "openai" and organization else _CapturedSDKHeader() + return { + "OpenAI-Organization": request_organization, + "OpenAI-Project": _CapturedSDKHeader(), + } + + +def _vertex_project_from_environment(): + # Google Auth treats an empty GOOGLE_CLOUD_PROJECT as masking GCLOUD_PROJECT. + return os.environ.get("VERTEXAI_PROJECT") or os.environ.get( + "GOOGLE_CLOUD_PROJECT", os.environ.get("GCLOUD_PROJECT"), + ) + + +def _get_bedrock_model_region(model: str, model_id=None) -> str | None: + """Resolve only the region carried by the request, without ambient AWS discovery.""" + from litellm.llms.bedrock.common_utils import BedrockModelInfo + + # LiteLLM 1.100.0 consumes model_id before Invoke resolves its region. + # Only Converse uses that separate ID or region/model path for routing. + is_converse = BedrockModelInfo.get_bedrock_route(model) == "converse" + if not is_converse: + model_id = None + candidate = model_id or model + if not isinstance(candidate, str): + return None + if not model_id: + for prefix in ("bedrock/converse/", "bedrock/", "converse/"): + if candidate.startswith(prefix): + candidate = candidate[len(prefix):] + break + arn_candidate = candidate + if not model_id: + arn_candidate = arn_candidate.removeprefix("invoke/") + # Match the native chat model prefixes in LiteLLM 1.100.0 without + # decoding an ARN or changing Converse's region/model path grammar. + for prefix in ("llama/", "deepseek_r1/", "openai/", "qwen2/", "qwen3/", "moonshot/", "nova-2/", "nova/"): + if arn_candidate.startswith(prefix): + arn_candidate = arn_candidate[len(prefix):] + break + arn = re.match(r"\Aarn:aws(?:-[a-z0-9-]+)?:bedrock:([a-z0-9-]+):", arn_candidate) + if arn: + return arn.group(1) + # Native Converse only recognizes region/model paths when model_id is absent. + if is_converse and not model_id: + region, separator, _ = candidate.partition("/") + if separator and region in litellm.AmazonBedrockGlobalConfig().get_all_regions(): + return region + return None + + +def _guard_request_routing_globals(provider: str | None, params: dict) -> dict: + """Reject process-wide LiteLLM routing fallbacks for one request.""" + if getattr(litellm, "api_base", None) and ( + "api_base" not in params or provider in LITELLM_GLOBAL_FIRST_API_BASE_PROVIDERS + ): + raise ValueError(f"Refusing process-wide LiteLLM API base fallback for provider {provider or 'unknown'}") + if provider == "gdc" and not params.get("api_base") and getattr(litellm, "gdc_api_base", None): + raise ValueError("Refusing process-wide LiteLLM API base fallback for provider gdc") + if provider == "azure" and "api_version" not in params and getattr(litellm, "api_version", None): + raise ValueError("Refusing process-wide LiteLLM API version fallback for provider azure") + organization_is_guarded = any( + header.lower() == "openai-organization" + for header in (params.get("headers") or {}) + ) + if ( + provider == "openai" + and "organization" not in params + and not organization_is_guarded + and getattr(litellm, "organization", None) + ): + raise ValueError("Refusing process-wide LiteLLM organization fallback for provider openai") + # GDC's complete project URL owns its routing; host-only URLs otherwise + # inherit the same process-wide project/location fallbacks as Vertex. + if provider == "vertex_ai" or (provider == "gdc" and "/v1/projects/" not in (params.get("api_base") or "")): + for parameter, global_name in ( + ("vertex_project", "vertex_project"), + ("vertex_location", "vertex_location"), + ): + if parameter not in params and getattr(litellm, global_name, None): + raise ValueError(f"Refusing process-wide LiteLLM {parameter} fallback for provider {provider}") + routing_environment_variables = dict(PROVIDER_ROUTING_ENV_VARS.get(provider, {})) + api_base_environment_variables = list(PROVIDER_API_BASE_ENV_VARS.get(provider, ())) + provider_config = JSONProviderRegistry.get(provider) + api_base_env = getattr(provider_config, "api_base_env", None) + if api_base_env: + api_base_environment_variables.append(api_base_env) + if api_base_environment_variables: + routing_environment_variables.setdefault("api_base", tuple(api_base_environment_variables)) + for parameter, environment_variables in routing_environment_variables.items(): + live_value = next( + (os.environ.get(variable) for variable in environment_variables if os.environ.get(variable)), + None, + ) + if provider == "vertex_ai" and parameter == "vertex_project": + live_value = _vertex_project_from_environment() + if parameter not in params and live_value: + raise ValueError(f"Refusing live {parameter} environment fallback for provider {provider}") + if ( + provider == "chatgpt" + and parameter == "api_base" + and parameter in params + and live_value != params[parameter] + ): + # LiteLLM 1.100.0's ChatGPT chat transformation ignores the request + # api_base and re-reads these variables when resolving provider info. + raise ValueError("Refusing changed live api_base environment for provider chatgpt") + return params + + +def _has_provider_api_key_global(provider: str) -> bool: + """Return whether LiteLLM retains a process-wide key for this provider.""" + return any(getattr(litellm, name, None) for name in PROVIDER_API_KEY_GLOBALS.get(provider, ())) + + +def _has_live_provider_api_key_environment(provider: str) -> bool: + """Return whether a provider key is currently available in the environment.""" + environment_variables = PROVIDER_API_KEY_ENV_VARS.get(provider, ()) + if not environment_variables: + provider_config = JSONProviderRegistry.get(provider) + api_key_env = getattr(provider_config, "api_key_env", None) + environment_variables = (api_key_env,) if api_key_env else () + return any(os.environ.get(environment_variable) for environment_variable in environment_variables) + + +AWS_REQUEST_CREDENTIAL_KEYS = ( + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_region_name", +) + +LITELLM_AWS_CREDENTIAL_SELECTOR_ENV_VARS = ( + "AWS_PROFILE_NAME", + "AWS_ROLE_NAME", +) + +AWS_REQUEST_ENDPOINT_ENV_VARS = ( + "AWS_ENDPOINT_URL", + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME", + "AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME", +) + +AWS_CREDENTIAL_CHAIN_ENV_VARS = ( + *LITELLM_AWS_CREDENTIAL_SELECTOR_ENV_VARS, + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_CONFIG_FILE", + "AWS_SHARED_CREDENTIALS_FILE", + "BOTO_CONFIG", + "AWS_CREDENTIAL_FILE", + "AWS_ROLE_ARN", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_ROLE_SESSION_NAME", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_AUTHORIZATION_TOKEN", + "AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE", + "AWS_EC2_METADATA_DISABLED", + "AWS_EC2_METADATA_SERVICE_ENDPOINT", + "AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE", + "AWS_EC2_METADATA_V1_DISABLED", + "AWS_IMDS_USE_IPV6", + *AWS_REQUEST_ENDPOINT_ENV_VARS, + "AWS_ENDPOINT_URL_STS", + "AWS_STS_REGIONAL_ENDPOINTS", + "AWS_SECURITY_TOKEN", + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_DEFAULT_REGION", +) + +BEDROCK_MANTLE_REQUEST_CONTEXT_KEYS = ( + *AWS_REQUEST_CREDENTIAL_KEYS, + "aws_bedrock_runtime_endpoint", +) + +BEDROCK_MANTLE_REQUEST_BODY_EXCLUDED_KEYS = BEDROCK_MANTLE_REQUEST_CONTEXT_KEYS + +AWS_REQUEST_PROVIDERS = { + "bedrock", + "bedrock_mantle", + "sagemaker", + "sagemaker_chat", + "sagemaker_nova", +} + +LITELLM_GLOBAL_FIRST_API_BASE_PROVIDERS = { + "custom", + "gradient_ai", + "ollama", + "ollama_chat", + "triton", +} + +_bedrock_mantle_request_credentials = ContextVar("bedrock_mantle_request_credentials", default=None) +_bedrock_mantle_block_bearer = ContextVar("bedrock_mantle_block_bearer", default=False) +_BEDROCK_MANTLE_ORIGINAL_SIGNER = "_pr_agent_original_sign_request" +_BEDROCK_MANTLE_ORIGINAL_TOKEN_RESOLVER = "_pr_agent_original_resolve_bearer_token" +_bedrock_mantle_sign_request = None +_bedrock_mantle_resolve_bearer_token = None +if BedrockMantleAuthMixin is not None: + _bedrock_mantle_sign_request = getattr( + BedrockMantleAuthMixin.sign_request, + _BEDROCK_MANTLE_ORIGINAL_SIGNER, + BedrockMantleAuthMixin.sign_request, + ) + _bedrock_mantle_resolve_bearer_token = getattr( + BedrockMantleAuthMixin._resolve_bearer_token, + _BEDROCK_MANTLE_ORIGINAL_TOKEN_RESOLVER, + BedrockMantleAuthMixin._resolve_bearer_token, + ) + +_vertex_request_credentials = ContextVar("vertex_request_credentials", default=None) +_vertex_request_active = ContextVar("vertex_request_active", default=False) +_vertex_request_aws_environment = ContextVar("vertex_request_aws_environment", default=None) +_vertex_request_default_adc = ContextVar("vertex_request_default_adc", default=None) +_databricks_request_keyless = ContextVar("databricks_request_keyless", default=False) +_raw_api_key_guard_provider = ContextVar("raw_api_key_guard_provider", default=None) +_raw_api_key_guard_auth = ContextVar("raw_api_key_guard_auth", default=None) + + +def _raw_guard_has_header_only_auth(provider, params, headers): + """Recognize only initial native header-only authentication choices.""" + snapshot = _raw_api_key_guard_auth.get() + if not snapshot or snapshot["generic_key"]: + return False + if provider == "ragflow": + return True + if provider == "xai": + return not snapshot["xai_key"] and not params.get("use_xai_oauth") + if provider != "azure_ai" or any(name.lower() == "api-key" for name in headers): + return False + if ( + snapshot["azure_key"] or snapshot["azure_ad_token"] or snapshot["azure_refresh"] + or params.get("azure_ad_token") or params.get("azure_ad_token_provider") is not None + ): + return False + environment = snapshot["azure_environment"] + tenant = params.get("tenant_id") or environment["AZURE_TENANT_ID"] + client = params.get("client_id") or environment["AZURE_CLIENT_ID"] + secret = params.get("client_secret") or environment["AZURE_CLIENT_SECRET"] + username = params.get("azure_username") or environment["AZURE_USERNAME"] + password = params.get("azure_password") or environment["AZURE_PASSWORD"] + return not ((tenant and client and secret) or (username and password and client)) + + +def _install_raw_api_key_guard_override_bridge(provider): + """Undo generated auth headers without reopening a native key resolver.""" + if provider not in ("azure_ai", "ragflow", "xai"): + return + from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig + from litellm.llms.ragflow.chat.transformation import RAGFlowConfig + from litellm.llms.xai.chat.transformation import XAIChatConfig + + config = {"azure_ai": AzureAIStudioConfig, "ragflow": RAGFlowConfig, "xai": XAIChatConfig}[provider] + original = config.validate_environment + marker = "_pr_agent_original_raw_override_validate_environment" + if getattr(original, marker, None) is not None: + return + signature = inspect.signature(original) + if not {"api_key", "headers", "litellm_params"}.issubset(signature.parameters): + raise RuntimeError("LiteLLM's native HTTP authentication interface is incompatible") + + def validate_environment(self, *args, **kwargs): + if _raw_api_key_guard_provider.get() != provider or type(self) is not config: + return original(self, *args, **kwargs) + bound = signature.bind(self, *args, **kwargs) + headers = bound.arguments.get("headers") or {} + native_params = bound.arguments.get("litellm_params") or {} + companion = _azure_ad_responses_request.get() + if ( + provider == "azure_ai" and companion is not None and companion.get("companion") + and companion["generated_guard"] + and bound.arguments.get("api_key") == DUMMY_LITELLM_API_KEY + and native_params.get("custom_llm_provider") == provider + and native_params.get("azure_ad_token") == companion["dispatch_token"] + ): + # The outer validator selects AD only without a key. Keep the guard + # in the inner validator, whose captured resolver bypasses fallback. + bound.arguments["headers"] = dict(headers) + bound.arguments["api_key"] = None + bound.arguments["litellm_params"] = {**native_params, "api_key": DUMMY_LITELLM_API_KEY} + return original(*bound.args, **bound.kwargs) + authorization = [name for name in headers if name.lower() == "authorization"] + if ( + bound.arguments.get("api_key") != DUMMY_LITELLM_API_KEY + or native_params.get("custom_llm_provider") != provider + or len(authorization) != 1 + or not _raw_guard_has_header_only_auth(provider, native_params, headers) + ): + return original(self, *args, **kwargs) + name = authorization[0] + value = headers[name] + bound.arguments["headers"] = dict(headers) + # Keep the guard during validation, including nested native resolvers. + # Preserve non-auth effects such as Ragflow's model extraction. + result = original(*bound.args, **bound.kwargs) + for header in list(result): + if header.lower() == "authorization" or (provider == "azure_ai" and header.lower() == "api-key"): + del result[header] + result[name] = value + return result + + setattr(validate_environment, marker, original) + config.validate_environment = validate_environment + + +def _install_raw_api_key_guard_bridge(): + """Keep generated fallback guards from replacing explicit native HTTP auth.""" + from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig + + original = OpenAIGPTConfig.validate_environment + marker = "_pr_agent_original_raw_validate_environment" + if getattr(original, marker, None) is not None: + return + signature = inspect.signature(original) + if not {"api_key", "headers", "litellm_params"}.issubset(signature.parameters): + raise RuntimeError("LiteLLM's native HTTP authentication interface is incompatible") + + def validate_environment(self, *args, **kwargs): + provider = _raw_api_key_guard_provider.get() + # Overrides that call super() may have their own authentication policy. + if provider is None or type(self).validate_environment is not validate_environment: + return original(self, *args, **kwargs) + bound = signature.bind(self, *args, **kwargs) + headers = bound.arguments.get("headers") or {} + native_params = bound.arguments.get("litellm_params") or {} + if ( + native_params.get("custom_llm_provider") == provider + and bound.arguments.get("api_key") == DUMMY_LITELLM_API_KEY + and any(name.lower() == "authorization" for name in headers) + ): + # Only this pure validator sees None. Keep the guard in dispatch + # kwargs so other native consumers cannot discover a foreign key. + bound.arguments["api_key"] = None + return original(*bound.args, **bound.kwargs) + + setattr(validate_environment, marker, original) + OpenAIGPTConfig.validate_environment = validate_environment + + +def _install_databricks_keyless_bridge(): + """Restore native auth only for a placeholder generated by the active handler.""" + try: + from litellm.llms.databricks.common_utils import DatabricksBase + except ImportError as error: + raise RuntimeError("LiteLLM's Databricks authentication interface is unavailable") from error + original = getattr(DatabricksBase, "databricks_validate_environment", None) + marker = "_pr_agent_original_databricks_validate_environment" + if getattr(original, marker, None) is not None: + return + try: + signature = inspect.signature(original) + parameter = signature.parameters["api_key"] + if parameter.kind not in (parameter.POSITIONAL_OR_KEYWORD, parameter.KEYWORD_ONLY): + raise ValueError("Unsupported api_key parameter") + except (KeyError, TypeError, ValueError) as error: + raise RuntimeError("LiteLLM's Databricks authentication interface is incompatible") from error + + def validate_environment(self, *args, **kwargs): + if not _databricks_request_keyless.get(): + return original(self, *args, **kwargs) + bound = signature.bind(self, *args, **kwargs) + if bound.arguments.get("api_key") == DUMMY_LITELLM_API_KEY: + bound.arguments["api_key"] = None + # Keep native OAuth/SDK authentication; do not let the global-key guard + # overwrite its Authorization header. Their deployment configuration + # is not a request-level PR-Agent credential setting. + return original(*bound.args, **bound.kwargs) + + setattr(validate_environment, marker, original) + DatabricksBase.databricks_validate_environment = validate_environment + + +def _snapshot_cloud_sdk_project(directory): + """Capture Cloud SDK 583's project property without running a subprocess.""" + if "CLOUDSDK_CORE_PROJECT" in os.environ: + return os.environ["CLOUDSDK_CORE_PROJECT"].strip() or None + paths = [] + executable = shutil.which("gcloud.cmd" if os.name == "nt" else "gcloud") + if executable: + sdk_root = os.path.dirname(os.path.dirname(os.path.realpath(executable))) + if os.path.isdir(os.path.join(sdk_root, ".install")): + paths.append(os.path.join(sdk_root, "properties")) + configuration = os.environ.get("CLOUDSDK_ACTIVE_CONFIG_NAME") + migrate_legacy = False + if not configuration: + try: + with open(os.path.join(directory, "active_config"), encoding="utf-8") as active: + configuration = active.read() + except FileNotFoundError: + configuration = None + migrate_legacy = True + if configuration and configuration != "NONE" and not re.match(r"^[a-z][-a-z0-9]*$", configuration): + # The SDK removes an invalid activator before checking legacy files. + migrate_legacy = True + configuration = None + if not configuration: + legacy_path = os.path.join(directory, "properties") + legacy_contents = "" + if migrate_legacy: + try: + with open(legacy_path, encoding="utf-8") as legacy: + legacy_contents = legacy.read() + except FileNotFoundError: + pass + deprecated = ( + "# This properties file has been superseded by named configurations.\n" + "# Editing it will have no effect.\n\n" + ) + if legacy_contents and not legacy_contents.startswith(deprecated): + paths.append(legacy_path) + else: + paths.append(os.path.join(directory, "configurations", "config_default")) + elif configuration != "NONE": + paths.append(os.path.join(directory, "configurations", f"config_{configuration}")) + # Match the SDK's terminal-newline semantics without trimming the filename. + if configuration and configuration != "NONE" and not re.match(r"^[a-z][-a-z0-9]*$", configuration): + raise ValueError("Invalid Cloud SDK configuration name") + project = None + for path in paths: + parser = configparser.ConfigParser() + try: + with open(path, encoding="utf-8") as config: + parser.read_file(config) + except FileNotFoundError: + continue + project = parser.get("core", "project", fallback=project) + return (project.strip() or None) if project is not None else None + + +def _load_vertex_default_adc(snapshot, project_id): + """Load only the captured ADC source, retaining native credential refresh.""" + from google.auth import _default, exceptions + from google.auth import credentials as google_credentials + from google.auth.transport.requests import Request + + if snapshot["error"]: + raise ValueError(f"Unable to snapshot default Vertex credentials: {snapshot['error']}") + info = snapshot["info"] + scopes = ["https://www.googleapis.com/auth/cloud-platform"] + credential_project = None + if info is None: + # Skip file-based discovery permanently when no file was captured. + # Retain google-auth 2.55.1's managed-runtime discovery and refresh. + for name in ("_get_gae_credentials", "_get_gce_credentials"): + if not callable(getattr(_default, name, None)): + raise RuntimeError("Google Auth managed ADC interface is unavailable") + loaded, credential_project = (None, None) + if snapshot["gae_runtime"] == "python27": + loaded, credential_project = _default._get_gae_credentials() + if loaded is None: + loaded, credential_project = _default._get_gce_credentials(request=Request()) + if loaded is None: + raise exceptions.DefaultCredentialsError("No captured file or managed-runtime Vertex credentials available") + else: + credential_type = info.get("type") + if credential_type == "authorized_user": + from google.oauth2.credentials import Credentials + + loaded = Credentials.from_authorized_user_info(info, scopes=scopes) + elif credential_type == "service_account": + from google.oauth2.service_account import Credentials + + loaded = Credentials.from_service_account_info(info, scopes=scopes) + credential_project = loaded.project_id + elif credential_type == "impersonated_service_account": + from google.auth.impersonated_credentials import Credentials + + loaded = Credentials.from_impersonated_service_account_info(info, scopes=scopes) + elif credential_type == "external_account_authorized_user": + from google.auth.external_account_authorized_user import Credentials + + loaded = Credentials.from_info(info) + elif credential_type == "gdch_service_account": + from google.oauth2.gdch_credentials import ServiceAccountCredentials + + loaded = ServiceAccountCredentials.from_service_account_info(info) + credential_project = info.get("project") + elif credential_type == "external_account": + source = info.get("credential_source", {}) + if not isinstance(source, dict): + raise ValueError("Invalid Vertex credential source") + if "executable" in source: + raise ValueError("Vertex executable credentials are incompatible with request isolation") + if info.get("subject_token_type") == "urn:ietf:params:aws:token-type:aws4_request": + loaded = _vertex_aws_credentials_from_snapshot(info, snapshot["aws_environment"], scopes) + else: + from google.auth.identity_pool import Credentials + + loaded = Credentials.from_info(info, scopes=scopes) + else: + raise ValueError("Unsupported default Vertex credential type") + loaded = google_credentials.with_scopes_if_required(loaded, scopes) + quota_project = snapshot["quota_project"] + if isinstance(loaded, google_credentials.CredentialsWithQuotaProject): + # Apply only the captured environment override; never consult a later one. + if quota_project or info is None: + loaded = loaded.with_quota_project(quota_project) + if ( + info is not None and info.get("type") == "external_account" + and not project_id and not snapshot["sdk_project"] + ): + credential_project = loaded.get_project_id(request=Request()) + resolved_project = project_id or credential_project or snapshot["sdk_project"] + if not resolved_project: + if snapshot["sdk_project_error"]: + raise ValueError(f"Unable to snapshot Cloud SDK project: {snapshot['sdk_project_error']}") + raise ValueError("Could not resolve project_id") + if not isinstance(resolved_project, str): + raise TypeError("Expected project_id to be a str") + return loaded, resolved_project + + +def _install_vertex_default_adc_bridge(): + """Recognize handler-owned ADC cache keys without changing unrelated LiteLLM callers.""" + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + original = VertexBase.load_auth + marker = "_pr_agent_original_default_adc_load_auth" + if getattr(original, marker, None) is not None: + return + + def load_auth(self, credentials, project_id): + snapshot = _vertex_request_default_adc.get() + if snapshot is None or credentials != snapshot["cache_key"]: + return original(self, credentials, project_id) + loaded, resolved_project = _load_vertex_default_adc(snapshot, project_id) + self.refresh_auth(loaded) + return loaded, resolved_project + + load_auth.__dict__.update(vars(original)) + setattr(load_auth, marker, original) + VertexBase.load_auth = load_auth + + +def _vertex_aws_credentials_from_snapshot(info, environment, scopes): + """Bind the native AWS WIF source without replacing Google's signer or refresh flow.""" + from google.auth import aws, exceptions + + environment = dict(environment) + source = dict(info["credential_source"]) + + class SnapshotSupplier(aws.AwsSecurityCredentialsSupplier): + def _metadata(self, request, url, *, method="GET", headers=None): + if not url: + raise exceptions.RefreshError("Missing AWS WIF metadata endpoint") + response = request(url=url, method=method, headers=headers) + if response.status != 200: + raise exceptions.RefreshError("Unable to retrieve AWS WIF metadata") + return response.data.decode("utf-8") if isinstance(response.data, bytes) else response.data + + def _metadata_headers(self, request): + token_url = source.get("imdsv2_session_token_url") + if token_url is None: + return None + token = self._metadata( + request, token_url, method="PUT", headers={"X-aws-ec2-metadata-token-ttl-seconds": "300"}, + ) + return {"X-aws-ec2-metadata-token": token} + + def get_aws_security_credentials(self, context, request): + access_key = environment.get("AWS_ACCESS_KEY_ID") + secret_key = environment.get("AWS_SECRET_ACCESS_KEY") + if access_key and secret_key: + return aws.AwsSecurityCredentials(access_key, secret_key, environment.get("AWS_SESSION_TOKEN")) + headers = self._metadata_headers(request) + url = source.get("url") + role = self._metadata(request, url, headers=headers) + credentials = json.loads(self._metadata(request, f"{url}/{role}", headers=headers)) + return aws.AwsSecurityCredentials( + credentials.get("AccessKeyId"), credentials.get("SecretAccessKey"), credentials.get("Token"), + ) + + def get_aws_region(self, context, request): + for variable in ("AWS_REGION", "AWS_DEFAULT_REGION"): + region = environment.get(variable) + if region is not None: + return region + return self._metadata(request, source.get("region_url"), headers=self._metadata_headers(request))[:-1] + + supplier = SnapshotSupplier() + + class SnapshotCredentials(aws.Credentials): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # google-auth 2.55.1 cannot combine a supplier with credential_source + # through its public constructor. Preserve native validation and the + # verification URL; reapply the supplier on copies and impersonation. + if not isinstance( + getattr(self, "_aws_security_credentials_supplier", None), aws.AwsSecurityCredentialsSupplier, + ): + raise RuntimeError("Google Auth AWS WIF supplier interface is unavailable") + self._aws_security_credentials_supplier = supplier + + return SnapshotCredentials.from_info(info, scopes=scopes) + + +def _install_vertex_executable_guard(): + """Reject executable ADC without replacing LiteLLM's native token cache.""" + from google.auth import pluggable + + credentials = pluggable.Credentials + marker = "_pr_agent_original_vertex_executable" + replacements = {} + + def guard(original): + def guarded(*args, **kwargs): + if _vertex_request_active.get(): + raise ValueError("Vertex executable credentials are incompatible with request isolation") + return original(*args, **kwargs) + + setattr(guarded, marker, original) + return guarded + + # LiteLLM 1.100.0 reads expired (sync) or token_state (async) on the + # actual cached credential before returning its token. Guarding the + # descriptors also covers cache replacement during an async lock wait. + for name in ("from_info", "refresh", "expired", "token_state"): + descriptor = inspect.getattr_static(credentials, name, None) + if isinstance(descriptor, classmethod): + original = descriptor.__func__ + elif isinstance(descriptor, property): + original = descriptor.fget + else: + original = descriptor + _require_litellm_interface(original, f"Google executable credentials {name}") + if getattr(original, marker, None) is not None: + continue + wrapped = guard(original) + if isinstance(descriptor, classmethod): + wrapped = classmethod(wrapped) + elif isinstance(descriptor, property): + wrapped = property(wrapped, descriptor.fset, descriptor.fdel, descriptor.__doc__) + replacements[name] = wrapped + + # Validate every entry point before installing any of the guards. + for name, replacement in replacements.items(): + setattr(credentials, name, replacement) + + +def _install_vertex_wif_project_bridge(): + """Restore ADC project discovery for request-owned WIF JSON in LiteLLM 1.100.0.""" + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + original = VertexBase.load_auth + marker = "_pr_agent_original_vertex_wif_load_auth" + if getattr(original, marker, None) is not None: + return + + def load_auth(self, credentials, project_id): + snapshot = _vertex_request_credentials.get() + if not isinstance(snapshot, dict) or snapshot.get("type") != "external_account": + return original(self, credentials, project_id) + try: + json_obj = json.loads(credentials) if isinstance(credentials, str) else credentials + except (ValueError, TypeError): + return original(self, credentials, project_id) + if json_obj != snapshot: + return original(self, credentials, project_id) + + # Preserve LiteLLM 1.100.0's WIF factories, including its explicit AWS + # supplier. Only project discovery is missing from its JSON load path. + scopes = ["https://www.googleapis.com/auth/cloud-platform"] + source = json_obj.get("credential_source", {}) + environment_id = source.get("environment_id", "") if isinstance(source, dict) else "" + aws_environment = _vertex_request_aws_environment.get() + bind_aws = isinstance(environment_id, str) and "aws" in environment_id and aws_environment is not None + if project_id is not None and not bind_aws: + return original(self, credentials, project_id) + if isinstance(environment_id, str) and "aws" in environment_id: + from litellm.llms.vertex_ai.vertex_ai_aws_wif import VertexAIAwsWifAuth + + aws_params = VertexAIAwsWifAuth.extract_aws_params(json_obj) + if aws_params: + loaded = VertexAIAwsWifAuth.credentials_from_explicit_aws(json_obj, aws_params, scopes) + elif bind_aws: + loaded = _vertex_aws_credentials_from_snapshot(json_obj, aws_environment, scopes) + else: + loaded = self._credentials_from_identity_pool_with_aws(json_obj, scopes) + elif isinstance(source, dict) and "executable" in source: + raise ValueError("Vertex executable credentials are incompatible with request isolation") + else: + loaded = self._credentials_from_identity_pool(json_obj, scopes) + + from google.auth.transport.requests import Request + + # Use the same credential for discovery, refresh, and the upstream cache; + # never reload ADC or apply a process-wide quota project override here. + resolved_project = project_id if project_id is not None else loaded.get_project_id(request=Request()) + self.refresh_auth(loaded) + if not resolved_project: + raise ValueError("Could not resolve project_id") + if not isinstance(resolved_project, str): + raise TypeError(f"Expected project_id to be a str but got {type(resolved_project)}") + return loaded, resolved_project + + load_auth.__dict__.update(vars(original)) + setattr(load_auth, marker, original) + VertexBase.load_auth = load_auth + + +def _install_vertex_impersonated_credentials_bridge(): + """Add LiteLLM 1.100.0's missing ADC type without replacing its refresh/cache path.""" + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + original = VertexBase._credentials_from_service_account + marker = "_pr_agent_original_vertex_service_account_loader" + if getattr(original, marker, None) is not None: + return + + def load_credentials(self, json_obj, scopes): + request_credentials = _vertex_request_credentials.get() + if ( + isinstance(json_obj, dict) + and json_obj.get("type") == "impersonated_service_account" + and request_credentials is not None + and json_obj == request_credentials + ): + from google.auth import impersonated_credentials + + # The generic ADC loader may apply a live GOOGLE_CLOUD_QUOTA_PROJECT. + # This constructor uses only the captured configuration, including quota. + return impersonated_credentials.Credentials.from_impersonated_service_account_info( + json_obj, scopes=scopes, + ) + return original(self, json_obj, scopes) + + setattr(load_credentials, marker, original) + VertexBase._credentials_from_service_account = load_credentials + + +_anthropic_request_auth_token = ContextVar("anthropic_request_auth_token", default=None) +_ANTHROPIC_ORIGINAL_API_KEY_RESOLVER = "_pr_agent_original_get_api_key" +_ANTHROPIC_ORIGINAL_AUTH_TOKEN_RESOLVER = "_pr_agent_original_get_auth_token" +_anthropic_get_api_key = getattr(AnthropicModelInfo, "get_api_key", None) +_anthropic_get_api_key = getattr( + _anthropic_get_api_key, + _ANTHROPIC_ORIGINAL_API_KEY_RESOLVER, + _anthropic_get_api_key, +) +_anthropic_get_auth_token = getattr(AnthropicModelInfo, "get_auth_token", None) +_anthropic_get_auth_token = getattr( + _anthropic_get_auth_token, + _ANTHROPIC_ORIGINAL_AUTH_TOKEN_RESOLVER, + _anthropic_get_auth_token, +) + + +def _resolve_anthropic_api_key(api_key=None): + """Turn PR-Agent's request-local guard back into Anthropic's keyless state.""" + request_auth = _anthropic_request_auth_token.get() + if request_auth is not None and request_auth["generated_guard"] and api_key == DUMMY_LITELLM_API_KEY: + return None + return _anthropic_get_api_key(api_key) + + +def _resolve_anthropic_auth_token(auth_token=None): + """Resolve only the bearer token captured for the active Anthropic request.""" + request_auth = _anthropic_request_auth_token.get() + if request_auth is not None: + return request_auth["auth_token"] + return _anthropic_get_auth_token(auth_token) + + +def _install_anthropic_auth_token_bridge(): + """Install request-local Anthropic credential resolvers without stacking wrappers.""" + global _anthropic_get_api_key, _anthropic_get_auth_token + _require_litellm_interface(AnthropicModelInfo, "AnthropicModelInfo", ("get_api_key", "get_auth_token")) + current_api_key_resolver = AnthropicModelInfo.get_api_key + _anthropic_get_api_key = getattr( + current_api_key_resolver, + _ANTHROPIC_ORIGINAL_API_KEY_RESOLVER, + current_api_key_resolver, + ) + setattr(_resolve_anthropic_api_key, _ANTHROPIC_ORIGINAL_API_KEY_RESOLVER, _anthropic_get_api_key) + AnthropicModelInfo.get_api_key = staticmethod(_resolve_anthropic_api_key) + current_auth_token_resolver = AnthropicModelInfo.get_auth_token + _anthropic_get_auth_token = getattr( + current_auth_token_resolver, + _ANTHROPIC_ORIGINAL_AUTH_TOKEN_RESOLVER, + current_auth_token_resolver, + ) + setattr( + _resolve_anthropic_auth_token, + _ANTHROPIC_ORIGINAL_AUTH_TOKEN_RESOLVER, + _anthropic_get_auth_token, + ) + AnthropicModelInfo.get_auth_token = staticmethod(_resolve_anthropic_auth_token) + + +def _resolve_bedrock_mantle_bearer_token(api_key): + """Prevent LiteLLM's generic key fallback while preserving request-local SigV4.""" + if _bedrock_mantle_block_bearer.get(): + return None if api_key == DUMMY_LITELLM_API_KEY else "" + return _bedrock_mantle_resolve_bearer_token(api_key) + + +def _sign_bedrock_mantle_request(self, *args, **kwargs): + """Bridge request-local AWS credentials to LiteLLM's Bedrock Mantle signer.""" + request_credentials = _bedrock_mantle_request_credentials.get() + if request_credentials is not None: + if "request_data" in kwargs: + request_data = kwargs["request_data"] + else: + try: + parameters = tuple(inspect.signature(_bedrock_mantle_sign_request).parameters.values()) + request_data_position = next( + index - 1 + for index, parameter in enumerate(parameters) + if parameter.name == "request_data" + and parameter.kind in (parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD) + ) + except (StopIteration, TypeError, ValueError): + request_data_position = None + if request_data_position is None or len(args) <= request_data_position: + raise RuntimeError( + "LiteLLM's Bedrock Mantle signer did not receive request_data; " + "request-local AWS credentials cannot be removed from the request body" + ) + request_data = args[request_data_position] + if not isinstance(request_data, dict): + raise RuntimeError( + "LiteLLM's Bedrock Mantle signer received invalid request_data; " + "request-local AWS credentials cannot be removed from the request body" + ) + for key in BEDROCK_MANTLE_REQUEST_BODY_EXCLUDED_KEYS: + request_data.pop(key, None) + if _bedrock_mantle_block_bearer.get(): + if "api_key" in kwargs: + kwargs["api_key"] = "" + else: + try: + parameters = tuple(inspect.signature(_bedrock_mantle_sign_request).parameters.values()) + api_key_parameter = next( + (index, parameter) + for index, parameter in enumerate(parameters) + if parameter.name == "api_key" + ) + except (StopIteration, TypeError, ValueError): + raise RuntimeError( + "LiteLLM's Bedrock Mantle signer did not expose api_key; " + "request-local bearer isolation cannot be applied" + ) + api_key_index, api_key_parameter = api_key_parameter + if api_key_parameter.kind == api_key_parameter.KEYWORD_ONLY: + kwargs["api_key"] = "" + elif api_key_parameter.kind in ( + api_key_parameter.POSITIONAL_ONLY, + api_key_parameter.POSITIONAL_OR_KEYWORD, + ): + api_key_position = api_key_index - 1 + if len(args) > api_key_position: + args = (*args[:api_key_position], "", *args[api_key_position + 1:]) + else: + kwargs["api_key"] = "" + else: + raise RuntimeError( + "LiteLLM's Bedrock Mantle signer has an unsupported api_key parameter; " + "request-local bearer isolation cannot be applied" + ) + if request_credentials: + if "optional_params" in kwargs: + kwargs["optional_params"] = {**(kwargs["optional_params"] or {}), **request_credentials} + else: + try: + parameters = tuple(inspect.signature(_bedrock_mantle_sign_request).parameters.values()) + optional_params_position = next( + index - 1 + for index, parameter in enumerate(parameters) + if parameter.name == "optional_params" + and parameter.kind in (parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD) + ) + except (StopIteration, TypeError, ValueError): + optional_params_position = None + if optional_params_position is None or len(args) <= optional_params_position: + raise RuntimeError( + "LiteLLM's Bedrock Mantle signer did not receive optional_params; " + "request-local AWS credentials cannot be applied" + ) + args = ( + *args[:optional_params_position], + {**(args[optional_params_position] or {}), **request_credentials}, + *args[optional_params_position + 1:], + ) + return _bedrock_mantle_sign_request(self, *args, **kwargs) + + +def _install_bedrock_mantle_signer_bridge(): + """Install the signer bridge without wrapping it again after a module reload.""" + global _bedrock_mantle_resolve_bearer_token, _bedrock_mantle_sign_request + if BedrockMantleAuthMixin is None: + return + current_signer = BedrockMantleAuthMixin.sign_request + _bedrock_mantle_sign_request = getattr( + current_signer, + _BEDROCK_MANTLE_ORIGINAL_SIGNER, + current_signer, + ) + setattr(_sign_bedrock_mantle_request, _BEDROCK_MANTLE_ORIGINAL_SIGNER, _bedrock_mantle_sign_request) + BedrockMantleAuthMixin.sign_request = _sign_bedrock_mantle_request + current_resolver = BedrockMantleAuthMixin._resolve_bearer_token + _bedrock_mantle_resolve_bearer_token = getattr( + current_resolver, + _BEDROCK_MANTLE_ORIGINAL_TOKEN_RESOLVER, + current_resolver, + ) + setattr( + _resolve_bedrock_mantle_bearer_token, + _BEDROCK_MANTLE_ORIGINAL_TOKEN_RESOLVER, + _bedrock_mantle_resolve_bearer_token, + ) + BedrockMantleAuthMixin._resolve_bearer_token = staticmethod(_resolve_bedrock_mantle_bearer_token) def _as_bool(value, default: bool) -> bool: @@ -47,6 +1893,15 @@ def _as_bool(value, default: bool) -> bool: return default +def _as_list(value) -> list: + """Parse a config value that may arrive as an list[str] (toml) or a string (env override).""" + if isinstance(value, (list, tuple)): + return [str(v).strip() for v in value if str(v).strip()] + if isinstance(value, str): + return [v.strip() for v in value.split(",") if v.strip()] + return [] + + def _configured_client_retries(): """config.num_retries as a non-negative int, or None (unset/invalid = client defaults). @@ -82,256 +1937,211 @@ def _should_retry_same_model(exc: BaseException) -> bool: class LiteLLMAIHandler(BaseAiHandler): - """ - This class handles interactions with the OpenAI API for chat completions. - It initializes the API key and other settings from a configuration file, - and provides a method for performing chat completions using the OpenAI ChatCompletion API. + """Handle chat completions across supported providers through LiteLLM. + + Request isolation covers PR-Agent settings and captured provider environment values, + not deployment-owned LiteLLM secret managers. Embedding applications must keep + process environment and LiteLLM globals stable while requests run; this handler + is not a sandbox for external code mutating shared authentication or routing state. """ def __init__(self): - """ - Initializes the OpenAI API key and other settings from a configuration file. - Raises a ValueError if the OpenAI key is missing. - """ - self.azure = False - self.api_base = None + """Initialize provider credentials and request settings from configuration.""" + _require_litellm_interface( + JSONProviderRegistry, "JSONProviderRegistry", ("list_providers", "get", "exists", "supports_responses_api"), + ) + settings = get_settings() + self._azure_ad = bool(settings.get("AZURE_AD.CLIENT_ID", None)) + self._azure_oidc_environment = {name: os.environ.get(name) for name in AZURE_OIDC_ENV_VARS} + self._azure_oidc_auth_environment = {name: os.environ.get(name) for name in AZURE_OIDC_AUTH_ENV_VARS} + self._azure_companion_auth = bool( + self._azure_oidc_environment["AZURE_CLIENT_ID"] and ( + (self._azure_oidc_environment["AZURE_TENANT_ID"] + and self._azure_oidc_auth_environment["AZURE_CLIENT_SECRET"]) + or (self._azure_oidc_auth_environment["AZURE_USERNAME"] + and self._azure_oidc_auth_environment["AZURE_PASSWORD"]) + ) + ) + self._raw_guard_auth_snapshot = { + "generic_key": bool(getattr(litellm, "api_key", None)), + "xai_key": bool(getattr(litellm, "xai_key", None)), + "azure_key": bool( + getattr(litellm, "azure_key", None) or os.environ.get("AZURE_OPENAI_API_KEY") + or os.environ.get("AZURE_API_KEY") or self._azure_ad + or (settings.get("OPENAI.API_TYPE", None) == "azure" and settings.get("OPENAI.KEY", None)) + ), + "azure_ad_token": bool(os.environ.get("AZURE_AD_TOKEN")), + "azure_refresh": getattr(litellm, "enable_azure_ad_token_refresh", False) is True, + "azure_environment": {**self._azure_oidc_environment, **self._azure_oidc_auth_environment}, + } + try: + self._azure_ad_credential = _get_azure_ad_credential(settings) if self._azure_ad else None + except Exception as e: + get_logger().error(f"Failed to create Azure AD credential: {type(e).__name__}") + raise + self.azure = settings.get("OPENAI.API_TYPE", None) == "azure" or self._azure_ad + self._openai_api_base_is_azure = settings.get("OPENAI.API_TYPE", None) == "azure" or ( + self._azure_ad and not settings.get("AZURE_AD.API_BASE", None) + ) self.repetition_penalty = None + self._aws_use_imds = False self._aws_imds_mode = False self._aws_static_creds = None + self._aws_environment_creds = None + self._aws_active_creds = {} + self._aws_environment_credentials_incomplete = False self._aws_imds_fell_back = False self._aws_boto3_creds = None # original boto3 credentials object for IMDS refresh + self._aws_region_name = None + self._aws_credential_chain_environment = {} + self._aws_credential_chain_files = {} self._aws_bedrock_lock = asyncio.Lock() + self._vertex_credentials, self._vertex_credentials_error = self._snapshot_vertex_credentials() + self._vertex_aws_environment = { + variable: os.environ.get(variable) + for variable in ( + "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "AWS_REGION", "AWS_DEFAULT_REGION", + ) + } + self._vertex_default_adc = None + self._vertex_gac_adc = None + if ( + self._vertex_credentials is not None + and not os.environ.get("VERTEXAI_CREDENTIALS") + and os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") + ): + # Non-SDK GAC uses Google's ADC factories, not LiteLLM's explicit + # Vertex JSON factories. Reuse the captured loader without SDK + # project discovery or later file/environment fallback. + snapshot = { + "source": "gac", "info": None, "error": None, + "sdk_project": None, "sdk_project_error": None, + "quota_project": os.environ.get("GOOGLE_CLOUD_QUOTA_PROJECT"), + "gae_runtime": None, "aws_environment": dict(self._vertex_aws_environment), + } + try: + snapshot["info"] = json.loads(self._vertex_credentials) + if not isinstance(snapshot["info"], dict): + raise ValueError("Invalid captured Vertex ADC") + except (TypeError, ValueError) as error: + snapshot["error"] = type(error).__name__ + identity = hashlib.sha256(json.dumps(snapshot, sort_keys=True).encode("utf-8")).hexdigest() + snapshot["cache_key"] = json.dumps({"_pr_agent_default_adc": identity}) + self._vertex_gac_adc = snapshot + if self._vertex_credentials is None and self._vertex_credentials_error is None: + from google.auth import _cloud_sdk - if get_settings().get("LITELLM.DISABLE_AIOHTTP", False): - litellm.disable_aiohttp_transport = True - if get_settings().get("OPENAI.KEY", None): - openai.api_key = get_settings().openai.key - litellm.openai_key = get_settings().openai.key - elif 'OPENAI_API_KEY' not in os.environ: - # These are process-wide globals and __init__ runs per request, so drop any - # key a previous request's provider branch left behind — chat_completion() - # forwards a truthy litellm.api_key, which would authenticate this request - # with (and leak) the earlier one's credential. The provider branches below - # repopulate it for the current request. - litellm.api_key = None - # Keyless OpenAI-compatible endpoints (vLLM, LM Studio, ...) still need *some* - # key or the OpenAI SDK raises "The api_key client option must be set". - # The placeholder must live in litellm.openai_key, which LiteLLM only consults - # on the OpenAI-compatible path: the global litellm.api_key is checked ahead of - # provider env vars (OPENROUTER_API_KEY, AZURE_API_KEY, ...) in LiteLLM's - # resolution chain, so a placeholder there silently shadows them. - litellm.openai_key = DUMMY_LITELLM_API_KEY - # Custom Bedrock endpoint (e.g. a VPC endpoint interface endpoint); litellm reads this - # directly from the environment regardless of the credentials model below - if not os.environ.get("AWS_BEDROCK_RUNTIME_ENDPOINT") and get_settings().get("aws.AWS_BEDROCK_RUNTIME_ENDPOINT"): - os.environ["AWS_BEDROCK_RUNTIME_ENDPOINT"] = get_settings().aws.AWS_BEDROCK_RUNTIME_ENDPOINT - if os.environ.get("AWS_USE_IMDS", "").strip().lower() in ("1", "true", "yes"): - import boto3 - import botocore.exceptions - session = boto3.Session() + snapshot = { + "info": None, "error": None, "sdk_project": None, "sdk_project_error": None, + "explicit_path": None, + "quota_project": os.environ.get("GOOGLE_CLOUD_QUOTA_PROJECT"), + "gae_runtime": os.environ.get("APPENGINE_RUNTIME"), + "aws_environment": dict(self._vertex_aws_environment), + } try: - creds = session.get_credentials() - if creds: - self._aws_boto3_creds = creds # store for refresh; avoids env-var re-read - self._write_frozen_aws_creds_to_env(creds.get_frozen_credentials()) - self._aws_imds_mode = True - get_logger().info("Using ambient AWS credentials from IMDS/task-role/IRSA") - else: - get_logger().warning( - "AWS_USE_IMDS is set but boto3 found no credentials; " - "falling through to static keys" - ) - except (botocore.exceptions.BotoCoreError, botocore.exceptions.ClientError): - # ClientError is intentionally not a BotoCoreError subclass in botocore's - # design; it is raised by STS-backed providers (AssumeRole, IRSA web-identity). - get_logger().exception( - "AWS_USE_IMDS: failed to resolve credentials via boto3; " - "falling through to static keys" - ) - if not os.environ.get("AWS_REGION_NAME"): - if get_settings().get("aws.AWS_REGION_NAME"): - os.environ["AWS_REGION_NAME"] = get_settings().aws.AWS_REGION_NAME + path = _cloud_sdk.get_application_default_credentials_path() + if os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") == path: + snapshot["explicit_path"] = path + try: + with open(path, encoding="utf-8") as adc_file: + snapshot["info"] = json.load(adc_file) + except FileNotFoundError: + pass else: + if not isinstance(snapshot["info"], dict): + raise ValueError("Invalid default Vertex credentials") try: - region = session.region_name - if region: - os.environ["AWS_REGION_NAME"] = region - get_logger().info(f"AWS region resolved from environment: {region}") - else: - get_logger().warning( - "AWS_USE_IMDS: could not determine AWS region; " - "set AWS_REGION_NAME explicitly" - ) - except Exception as e: - get_logger().warning(f"AWS_USE_IMDS: failed to resolve region via boto3: {e}") - if get_settings().get("aws.AWS_ACCESS_KEY_ID"): - if get_settings().aws.AWS_SECRET_ACCESS_KEY and get_settings().aws.AWS_REGION_NAME: - static_creds = { - "AWS_ACCESS_KEY_ID": get_settings().aws.AWS_ACCESS_KEY_ID, - "AWS_SECRET_ACCESS_KEY": get_settings().aws.AWS_SECRET_ACCESS_KEY, - "AWS_REGION_NAME": get_settings().aws.AWS_REGION_NAME, - } - static_token = get_settings().get("aws.AWS_SESSION_TOKEN", None) - if static_token: - static_creds["AWS_SESSION_TOKEN"] = static_token - if self._aws_imds_mode: - # IMDS succeeded; stash static keys for runtime fallback only - self._aws_static_creds = static_creds - else: - # IMDS failed; activate static credentials immediately and stash - # them so the runtime fallback path is also available if needed. - self._aws_static_creds = static_creds - os.environ["AWS_ACCESS_KEY_ID"] = static_creds["AWS_ACCESS_KEY_ID"] - os.environ["AWS_SECRET_ACCESS_KEY"] = static_creds["AWS_SECRET_ACCESS_KEY"] - os.environ["AWS_REGION_NAME"] = static_creds["AWS_REGION_NAME"] - if static_token: - os.environ["AWS_SESSION_TOKEN"] = static_token - elif "AWS_SESSION_TOKEN" in os.environ: - del os.environ["AWS_SESSION_TOKEN"] - get_logger().info( - "AWS_USE_IMDS: IMDS resolution failed; using static credentials" - ) - elif get_settings().get("aws.AWS_ACCESS_KEY_ID"): - assert get_settings().aws.AWS_SECRET_ACCESS_KEY and get_settings().aws.AWS_REGION_NAME, "AWS credentials are incomplete" - os.environ["AWS_ACCESS_KEY_ID"] = get_settings().aws.AWS_ACCESS_KEY_ID - os.environ["AWS_SECRET_ACCESS_KEY"] = get_settings().aws.AWS_SECRET_ACCESS_KEY - os.environ["AWS_REGION_NAME"] = get_settings().aws.AWS_REGION_NAME - static_token = get_settings().get("aws.AWS_SESSION_TOKEN", None) - if static_token: - os.environ["AWS_SESSION_TOKEN"] = static_token - elif "AWS_SESSION_TOKEN" in os.environ: - del os.environ["AWS_SESSION_TOKEN"] - if get_settings().get("LITELLM.DROP_PARAMS", None): - litellm.drop_params = get_settings().litellm.drop_params - if get_settings().get("LITELLM.SUCCESS_CALLBACK", None): - litellm.success_callback = get_settings().litellm.success_callback - if get_settings().get("LITELLM.FAILURE_CALLBACK", None): - litellm.failure_callback = get_settings().litellm.failure_callback - if get_settings().get("LITELLM.SERVICE_CALLBACK", None): - litellm.service_callback = get_settings().litellm.service_callback - # litellm's callbacks attach full prompt and response content — here, the whole - # PR diff — to whatever they emit, unless message logging is turned off. - if get_settings().get("LITELLM.TURN_OFF_MESSAGE_LOGGING", False): + snapshot["sdk_project"] = _snapshot_cloud_sdk_project(os.path.dirname(path)) + except (OSError, ValueError, configparser.Error) as error: + snapshot["sdk_project_error"] = type(error).__name__ + except (OSError, ValueError, configparser.Error) as error: + snapshot["error"] = type(error).__name__ + identity = hashlib.sha256(json.dumps(snapshot, sort_keys=True).encode("utf-8")).hexdigest() + snapshot["cache_key"] = json.dumps({"_pr_agent_default_adc": identity}) + self._vertex_default_adc = snapshot + if self._vertex_credentials: + try: + vertex_info = json.loads(self._vertex_credentials) + except (ValueError, TypeError): + vertex_info = None + if isinstance(vertex_info, dict) and vertex_info.get("type") == "external_account": + source = vertex_info.get("credential_source", {}) + environment_id = source.get("environment_id", "") if isinstance(source, dict) else "" + if isinstance(environment_id, str) and "aws" in environment_id: + # Separate different AWS identities without growing LiteLLM's + # cache for every new handler with the same captured source. + vertex_info["_pr_agent_aws_identity"] = hashlib.sha256( + json.dumps(self._vertex_aws_environment, sort_keys=True).encode("utf-8"), + ).hexdigest() + self._vertex_credentials = json.dumps(vertex_info) + self._provider_request_params = self._snapshot_provider_request_params(settings) + self._snowflake_account_id = os.environ.get("SNOWFLAKE_ACCOUNT_ID", "") + self._provider_environment_api_keys = self._snapshot_provider_environment_api_keys() + self._volcengine_ark_api_key = os.environ.get("ARK_API_KEY") + self._request_headers = self._snapshot_request_headers(settings) + sdk_headers = {} + for line in os.environ.get("OPENAI_CUSTOM_HEADERS", "").split("\n"): + if ":" not in line: + continue + name, value = line.split(":", 1) + sdk_headers[name.strip()] = value.strip() + _check_sdk_marker_collision(sdk_headers) + _check_sdk_marker_collision(self._request_headers) + self._sdk_header_defaults = { + "organization": os.environ.get("OPENAI_ORG_ID"), + "project": os.environ.get("OPENAI_PROJECT_ID"), + "custom_headers": sdk_headers, + } + openrouter_settings = settings.get("openrouter", {}) or {} + # Credentials and endpoints are isolated in _provider_request_params; keep this snapshot control-only. + self._openrouter_controls = { + key: copy.deepcopy(openrouter_settings.get(key)) + for key in ( + "provider_only", + "provider_order", + "allow_fallbacks", + "reasoning_effort", + "reasoning_max_tokens", + "max_tokens", + ) + } + self._default_reasoning_effort = getattr(settings.config, "reasoning_effort", None) + self._claude_thinking_controls = { + key: copy.deepcopy(settings.config.get(key, default)) + for key, default in ( + ("enable_claude_adaptive_thinking", False), + ("enable_claude_extended_thinking", False), + ("extended_thinking_budget_tokens", 2048), + ("extended_thinking_max_output_tokens", 4096), + ) + } + self._bedrock_model_id = settings.get("litellm.model_id", None) + self._custom_llm_provider = str( + getattr(settings.litellm, "custom_llm_provider", "") or "" + ).strip().lower() + self._anthropic_auth_token = os.environ.get("ANTHROPIC_AUTH_TOKEN") + self._request_provider_cache = {} + + if settings.get("LITELLM.DISABLE_AIOHTTP", False): + litellm.disable_aiohttp_transport = True + self._initialize_aws_request_credentials(settings) + if settings.get("LITELLM.DROP_PARAMS", None): + litellm.drop_params = settings.litellm.drop_params + if settings.get("LITELLM.SUCCESS_CALLBACK", None): + litellm.success_callback = settings.litellm.success_callback + if settings.get("LITELLM.FAILURE_CALLBACK", None): + litellm.failure_callback = settings.litellm.failure_callback + if settings.get("LITELLM.SERVICE_CALLBACK", None): + litellm.service_callback = settings.litellm.service_callback + # litellm callbacks attach full prompt and response content unless message logging is disabled. + if settings.get("LITELLM.TURN_OFF_MESSAGE_LOGGING", False): litellm.turn_off_message_logging = True - # With pr-agent's own telemetry enabled, its command span is the active parent and - # litellm's "otel" callback would skip its own request span, writing gen_ai attributes - # onto the command span instead. Keep the two layers separately aggregatable; - # setdefault leaves an explicit operator override in effect. - if self._litellm_otel_callback_enabled() and get_settings().get("OTEL.IS_ENABLED", False): + # Keep LiteLLM request spans separate from pr-agent command spans when both OTEL layers are enabled. + if self._litellm_otel_callback_enabled() and settings.get("OTEL.IS_ENABLED", False): os.environ.setdefault("USE_OTEL_LITELLM_REQUEST_SPAN", "true") - if get_settings().get("OPENAI.ORG", None): - litellm.organization = get_settings().openai.org - if get_settings().get("OPENAI.API_TYPE", None): - if get_settings().openai.api_type == "azure": - self.azure = True - litellm.azure_key = get_settings().openai.key - if get_settings().get("OPENAI.API_VERSION", None): - litellm.api_version = get_settings().openai.api_version - if get_settings().get("OPENAI.API_BASE", None): - litellm.api_base = get_settings().openai.api_base - self.api_base = get_settings().openai.api_base - if get_settings().get("ANTHROPIC.KEY", None): - litellm.anthropic_key = get_settings().anthropic.key - if get_settings().get("COHERE.KEY", None): - litellm.cohere_key = get_settings().cohere.key - if get_settings().get("GROQ.KEY", None): - litellm.api_key = get_settings().groq.key - if get_settings().get("SAMBANOVA.KEY", None): - litellm.api_key = get_settings().sambanova.key - if get_settings().get("REPLICATE.KEY", None): - litellm.replicate_key = get_settings().replicate.key - if get_settings().get("XAI.KEY", None): - litellm.api_key = get_settings().xai.key - if get_settings().get("HUGGINGFACE.KEY", None): - litellm.huggingface_key = get_settings().huggingface.key - if get_settings().get("HUGGINGFACE.API_BASE", None) and 'huggingface' in get_settings().config.model: - litellm.api_base = get_settings().huggingface.api_base - self.api_base = get_settings().huggingface.api_base - if get_settings().get("OLLAMA.API_BASE", None): - litellm.api_base = get_settings().ollama.api_base - self.api_base = get_settings().ollama.api_base - if get_settings().get("OLLAMA.API_KEY", None): - litellm.api_key = get_settings().ollama.api_key repetition_penalty = get_repetition_penalty() if repetition_penalty is not None: self.repetition_penalty = repetition_penalty - if get_settings().get("VERTEXAI.VERTEX_PROJECT", None): - litellm.vertex_project = get_settings().vertexai.vertex_project - litellm.vertex_location = get_settings().get( - "VERTEXAI.VERTEX_LOCATION", None - ) - # Google AI Studio - # SEE https://docs.litellm.ai/docs/providers/gemini - if get_settings().get("GOOGLE_AI_STUDIO.GEMINI_API_KEY", None): - os.environ["GEMINI_API_KEY"] = get_settings().google_ai_studio.gemini_api_key - - # Support deepseek models - if get_settings().get("DEEPSEEK.KEY", None): - os.environ['DEEPSEEK_API_KEY'] = get_settings().get("DEEPSEEK.KEY") - - # Support GLM (Z.AI / Zhipu) models - if get_settings().get("ZAI.KEY", None): - os.environ['ZAI_API_KEY'] = get_settings().get("ZAI.KEY") - - # Support Moonshot (Kimi) models - if get_settings().get("MOONSHOT.KEY", None): - os.environ['MOONSHOT_API_KEY'] = get_settings().get("MOONSHOT.KEY") - # Optional Moonshot endpoint override (e.g. China: https://api.moonshot.cn/v1) - if get_settings().get("MOONSHOT.API_BASE", None): - os.environ['MOONSHOT_API_BASE'] = get_settings().get("MOONSHOT.API_BASE") - - # Support Qwen (Alibaba DashScope) models - if get_settings().get("DASHSCOPE.KEY", None): - os.environ['DASHSCOPE_API_KEY'] = get_settings().get("DASHSCOPE.KEY") - - # Support Xiaomi MiMo models - if get_settings().get("XIAOMI_MIMO.KEY", None): - os.environ['XIAOMI_MIMO_API_KEY'] = get_settings().get("XIAOMI_MIMO.KEY") - - # Support deepinfra models - if get_settings().get("DEEPINFRA.KEY", None): - os.environ['DEEPINFRA_API_KEY'] = get_settings().get("DEEPINFRA.KEY") - - # Support mistral models - if get_settings().get("MISTRAL.KEY", None): - os.environ["MISTRAL_API_KEY"] = get_settings().get("MISTRAL.KEY") - - # Support codestral models - if get_settings().get("CODESTRAL.KEY", None): - os.environ["CODESTRAL_API_KEY"] = get_settings().get("CODESTRAL.KEY") - - # Support Databricks-hosted models (e.g. Azure Databricks serving endpoints). - # Uses PAT/key authentication via LiteLLM's env vars. - # SEE https://docs.litellm.ai/docs/providers/databricks - if get_settings().get("DATABRICKS.API_KEY", None): - os.environ["DATABRICKS_API_KEY"] = get_settings().get("DATABRICKS.API_KEY") - if get_settings().get("DATABRICKS.API_BASE", None): - os.environ["DATABRICKS_API_BASE"] = get_settings().get("DATABRICKS.API_BASE") - - # Check for Azure AD configuration - if get_settings().get("AZURE_AD.CLIENT_ID", None): - self.azure = True - # Generate access token using Azure AD credentials from settings - access_token = _get_azure_ad_token() - litellm.api_key = access_token - openai.api_key = access_token - - # Set API base from settings - self.api_base = get_settings().azure_ad.api_base - litellm.api_base = self.api_base - openai.api_base = self.api_base - - # Support for Openrouter models - if get_settings().get("OPENROUTER.KEY", None): - openrouter_api_key = get_settings().get("OPENROUTER.KEY", None) - os.environ["OPENROUTER_API_KEY"] = openrouter_api_key - litellm.api_key = openrouter_api_key - openai.api_key = openrouter_api_key - - openrouter_api_base = get_settings().get("OPENROUTER.API_BASE", "https://openrouter.ai/api/v1") - os.environ["OPENROUTER_API_BASE"] = openrouter_api_base - self.api_base = openrouter_api_base - litellm.api_base = openrouter_api_base # Models that only use user message self.user_message_only_models = USER_MESSAGE_ONLY_MODELS @@ -380,35 +2190,480 @@ def __init__(self): else: if raw_force_streaming_api_base_substrings: get_logger().warning( - "LITELLM.FORCE_STREAMING_API_BASE_SUBSTRINGS must be a list, tuple, or set. " - "Ignoring invalid value." + "LITELLM.FORCE_STREAMING_API_BASE_SUBSTRINGS must be a list, tuple, or set. " + "Ignoring invalid value." + ) + self.force_streaming_api_base_substrings = [] + + @staticmethod + def _litellm_otel_callback_enabled() -> bool: + """True when litellm's built-in OpenTelemetry callback is registered.""" + return any( + "otel" in (getattr(litellm, name, None) or []) + for name in ("callbacks", "success_callback", "failure_callback", "service_callback") + ) + + def _snapshot_provider_request_params(self, settings) -> dict: + """Capture provider credentials and endpoints for this handler instance.""" + provider_params = {} + for provider, setting_paths in PROVIDER_SETTING_PATHS.items(): + params = { + parameter: settings.get(setting_path, None) + for parameter, setting_path in setting_paths.items() + if settings.get(setting_path, None) + } + if params: + provider_params[provider] = params + + for provider, environment_variables in PROVIDER_API_BASE_ENV_VARS.items(): + if provider_params.get(provider, {}).get("api_base"): + continue + for environment_variable in environment_variables: + api_base = os.environ.get(environment_variable) + if api_base: + provider_params.setdefault(provider, {})["api_base"] = api_base + break + + if "api_base" not in provider_params.get("cloudflare", {}): + cloudflare_account_id = os.environ.get("CLOUDFLARE_ACCOUNT_ID") + if cloudflare_account_id: + provider_params.setdefault("cloudflare", {})["api_base"] = ( + f"https://api.cloudflare.com/client/v4/accounts/{cloudflare_account_id}/ai/v1" + ) + + for provider in ("watsonx", "watsonx_text"): + for parameter, environment_variables in PROVIDER_ROUTING_ENV_VARS[provider].items(): + if provider_params.get(provider, {}).get(parameter): + continue + for environment_variable in environment_variables: + value = os.environ.get(environment_variable) + if value: + provider_params.setdefault(provider, {})[parameter] = value + break + + aws_region = self._resolve_aws_region(settings) + if aws_region: + provider_params.setdefault("bedrock", {})["aws_region_name"] = aws_region + mantle_aws_region = aws_region + use_imds = os.environ.get("AWS_USE_IMDS", "").strip().lower() in ("1", "true", "yes") + has_static_credentials = all(settings.get(f"aws.{name}", None) for name in ( + "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME", + )) + if not use_imds and not has_static_credentials: + # Native Mantle ignores AWS_DEFAULT_REGION. Pin its default too, so + # the shared AWS credential snapshot cannot change the request region. + mantle_aws_region = ( + os.environ.get("AWS_REGION_NAME") + or settings.get("aws.AWS_REGION_NAME", None) + or os.environ.get("AWS_REGION") + or "us-east-1" + ) + bedrock_mantle_region = os.environ.get("BEDROCK_MANTLE_REGION") or mantle_aws_region + bedrock_mantle_api_base = provider_params.get("bedrock_mantle", {}).get("api_base") + if bedrock_mantle_api_base and MANTLE_HOST_RE is not None: + match = MANTLE_HOST_RE.match(bedrock_mantle_api_base.rstrip("/")) + if match: + bedrock_mantle_region = match.group(1) + if bedrock_mantle_region: + provider_params.setdefault("bedrock_mantle", {})["aws_region_name"] = bedrock_mantle_region + + for provider in JSONProviderRegistry.list_providers(): + provider_config = JSONProviderRegistry.get(provider) + if provider_config is None or provider_params.get(provider, {}).get("api_base"): + continue + api_base_env = getattr(provider_config, "api_base_env", None) + if not api_base_env: + continue + api_base = os.environ.get(api_base_env) + if api_base: + provider_params.setdefault(provider, {})["api_base"] = api_base + + openai_params = { + parameter: value + for parameter, value in { + "api_key": settings.get("OPENAI.KEY", None), + "api_base": ( + settings.get("OPENAI.API_BASE", None) + or os.environ.get("OPENAI_BASE_URL") + or os.environ.get("OPENAI_API_BASE") + ), + "api_version": settings.get("OPENAI.API_VERSION", None), + "organization": settings.get("OPENAI.ORG", None) or os.environ.get("OPENAI_ORGANIZATION"), + }.items() + if value + } + if openai_params: + provider_params["openai"] = openai_params + + bedrock_runtime_endpoint = ( + settings.get("aws.AWS_BEDROCK_RUNTIME_ENDPOINT", None) + or os.environ.get("AWS_BEDROCK_RUNTIME_ENDPOINT") + ) + if bedrock_runtime_endpoint: + for provider in ("bedrock", "bedrock_mantle"): + provider_params.setdefault(provider, {})[ + "aws_bedrock_runtime_endpoint" + ] = bedrock_runtime_endpoint + + configured_api_base = settings.get("OPENAI.API_BASE", None) + configured_api_version = settings.get("OPENAI.API_VERSION", None) + azure_api_base = os.environ.get("AZURE_API_BASE") + azure_api_version = os.environ.get("AZURE_API_VERSION") + azure_ad_api_base = settings.get("AZURE_AD.API_BASE", None) if self._azure_ad else None + if azure_ad_api_base: + request_api_base = azure_ad_api_base + request_api_version = configured_api_version or azure_api_version + elif self._openai_api_base_is_azure: + if configured_api_base: + request_api_base = configured_api_base + request_api_version = configured_api_version or azure_api_version + else: + request_api_base = azure_api_base + request_api_version = azure_api_version or configured_api_version + elif azure_api_base: + request_api_base = azure_api_base + request_api_version = azure_api_version or configured_api_version + else: + request_api_base = configured_api_base + request_api_version = configured_api_version or azure_api_version + + # The SDK endpoint alias is a fallback, not an override of configured routing. + request_api_base = request_api_base or os.environ.get("AZURE_OPENAI_ENDPOINT") + + # Cloudflare's key branch lets the SDK read this alias independently. + self._azure_sdk_ad_token = os.environ.get("AZURE_OPENAI_AD_TOKEN") + azure_params = { + parameter: value + for parameter, value in { + "api_key": ( + settings.get("OPENAI.KEY", None) + if settings.get("OPENAI.API_TYPE", None) == "azure" + else None + ), + "api_base": request_api_base, + "api_version": request_api_version, + "azure_ad_token": ( + os.environ.get("AZURE_AD_TOKEN") or self._azure_sdk_ad_token + ), + }.items() + if value + } + if azure_params: + provider_params["azure"] = {key: value for key, value in azure_params.items() if value} + + vertex_params = { + parameter: value + for parameter, value in { + "vertex_project": ( + settings.get("VERTEXAI.VERTEX_PROJECT", None) + or _vertex_project_from_environment() + ), + "vertex_location": ( + settings.get("VERTEXAI.VERTEX_LOCATION", None) + or os.environ.get("VERTEXAI_LOCATION") + or os.environ.get("VERTEX_LOCATION") + ), + }.items() + if value + } + vertex_credentials = self._vertex_credentials + if self._vertex_gac_adc is not None: + vertex_params["vertex_credentials"] = self._vertex_gac_adc["cache_key"] + elif vertex_credentials: + vertex_params["vertex_credentials"] = vertex_credentials + elif self._vertex_default_adc is not None: + vertex_params["vertex_credentials"] = self._vertex_default_adc["cache_key"] + if vertex_params: + provider_params.setdefault("vertex_ai", {}).update(vertex_params) + + if settings.get("OPENROUTER.KEY", None) or any( + os.environ.get(environment_variable) + for environment_variable in PROVIDER_API_KEY_ENV_VARS["openrouter"] + ): + openrouter_api_base = ( + settings.get("OPENROUTER.API_BASE", None) + or os.environ.get("OPENROUTER_API_BASE") + or "https://openrouter.ai/api/v1" + ) + provider_params.setdefault("openrouter", {}).setdefault("api_base", openrouter_api_base) + return provider_params + + @staticmethod + def _snapshot_vertex_credentials() -> tuple[str | None, str | None]: + """Capture explicit Vertex credentials before another request can change their environment.""" + vertex_credentials = os.environ.get("VERTEXAI_CREDENTIALS") + if vertex_credentials and not os.path.isfile(vertex_credentials): + try: + parsed_credentials = json.loads(vertex_credentials) + except (json.JSONDecodeError, TypeError): + pass + else: + if isinstance(parsed_credentials, dict): + return vertex_credentials, None + return None, "ValueError" + + credentials_path = vertex_credentials or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") + if not credentials_path: + return None, None + if not vertex_credentials: + from google.auth import _cloud_sdk + + # Match Google Auth's exact-path exception so the captured SDK + # resource project remains distinct from the billing quota project. + if credentials_path == _cloud_sdk.get_application_default_credentials_path(): + return None, None + try: + with open(credentials_path, encoding="utf-8") as credentials_file: + return credentials_file.read(), None + except OSError as e: + return None, type(e).__name__ + + @staticmethod + def _snapshot_provider_environment_api_keys() -> dict: + """Capture native provider API keys without mixing them with configured credentials.""" + provider_api_keys = {} + for provider, environment_variables in PROVIDER_API_KEY_ENV_VARS.items(): + for environment_variable in environment_variables: + api_key = os.environ.get(environment_variable) + if api_key: + provider_api_keys[provider] = api_key + break + for provider in JSONProviderRegistry.list_providers(): + provider_config = JSONProviderRegistry.get(provider) + if provider_config is None or provider in provider_api_keys: + continue + api_key_env = getattr(provider_config, "api_key_env", None) + if not api_key_env: + continue + api_key = os.environ.get(api_key_env) + if api_key: + provider_api_keys[provider] = api_key + return provider_api_keys + + def _captured_api_key(self, provider: str): + """Return the API key captured for a provider, preferring configured params over the environment.""" + return ( + getattr(self, "_provider_request_params", {}).get(provider, {}).get("api_key") + or getattr(self, "_provider_environment_api_keys", {}).get(provider) + ) + + @staticmethod + def _resolve_aws_region(settings) -> str | None: + """Resolve the AWS region from the environment and the configured aws settings.""" + return ( + os.environ.get("AWS_REGION_NAME") + or settings.get("aws.AWS_REGION_NAME", None) + or os.environ.get("AWS_REGION") + or os.environ.get("AWS_DEFAULT_REGION") + ) + + @staticmethod + def _snapshot_request_headers(settings) -> dict: + """Capture explicitly configured headers for every request from this handler.""" + raw_headers = settings.get("LITELLM.EXTRA_HEADERS", None) + if not raw_headers: + return {} + try: + request_headers = json.loads(raw_headers) + except (json.JSONDecodeError, TypeError) as e: + raise ValueError(f"LITELLM.EXTRA_HEADERS contains invalid JSON: {str(e)}") from e + if not isinstance(request_headers, dict): + raise ValueError("LITELLM.EXTRA_HEADERS must be a JSON object") + return request_headers + + def _initialize_aws_request_credentials(self, settings) -> None: + """Capture request-local credentials and synchronously resolve the opted-in AWS provider chain.""" + use_imds = os.environ.get("AWS_USE_IMDS", "").strip().lower() in ("1", "true", "yes") + self._aws_use_imds = use_imds + self._aws_credential_chain_environment = { + variable: os.environ.get(variable) + for variable in AWS_CREDENTIAL_CHAIN_ENV_VARS + } + request_region = self._resolve_aws_region(settings) + ambient_access_key = os.environ.get("AWS_ACCESS_KEY_ID") + ambient_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY") + if bool(ambient_access_key) != bool(ambient_secret_key): + self._aws_environment_credentials_incomplete = True + elif ambient_access_key and ambient_secret_key: + self._aws_environment_creds = { + "aws_access_key_id": ambient_access_key, + "aws_secret_access_key": ambient_secret_key, + # Prevent LiteLLM from falling back to a later ambient STS token. + # Only the opted-in boto3 chain recognizes the SECURITY alias; + # native LiteLLM's explicit keypair path uses SESSION alone. + "aws_session_token": ( + (os.environ.get("AWS_SECURITY_TOKEN") or os.environ.get("AWS_SESSION_TOKEN")) + if use_imds else os.environ.get("AWS_SESSION_TOKEN") + ) or "", + } + if request_region: + self._aws_environment_creds["aws_region_name"] = request_region + + static_access_key = settings.get("aws.AWS_ACCESS_KEY_ID", None) + if static_access_key: + static_secret_key = settings.get("aws.AWS_SECRET_ACCESS_KEY", None) + static_region = settings.get("aws.AWS_REGION_NAME", None) + if not (static_secret_key and static_region): + if not use_imds: + raise ValueError("AWS credentials are incomplete") + get_logger().warning( + "AWS_USE_IMDS is set but configured static AWS credentials are incomplete; " + "no static fallback is available" ) - self.force_streaming_api_base_substrings = [] + if static_secret_key and static_region: + self._aws_static_creds = { + "aws_access_key_id": static_access_key, + "aws_secret_access_key": static_secret_key, + # LiteLLM falls back to AWS_SESSION_TOKEN only when this value is None. + # An empty token keeps long-lived static keys isolated from ambient STS credentials. + "aws_session_token": settings.get("aws.AWS_SESSION_TOKEN", None) or "", + "aws_region_name": static_region, + } + + if not use_imds: + self._aws_active_creds = dict(self._aws_static_creds or self._aws_environment_creds or {}) + return + + if not (ambient_access_key or ambient_secret_key): + self._aws_credential_chain_files = self._snapshot_aws_credential_chain_files() + self._aws_region_name = request_region + self._initialize_aws_imds_credentials() + + def _initialize_aws_imds_credentials(self) -> bool: + """Resolve ambient AWS credentials during handler initialization without changing process credentials.""" + import boto3 + import botocore.exceptions + + self._validate_aws_credential_chain_environment() + region = self._aws_region_name + if self._aws_environment_credentials_incomplete: + if not self._aws_static_creds: + raise ValueError("AWS environment credentials are incomplete") + self._aws_active_creds = dict(self._aws_static_creds) + self._aws_imds_fell_back = True + get_logger().warning( + "AWS_USE_IMDS: ambient credentials are incomplete; using static credentials" + ) + return False + try: + session_kwargs = {} + if self._aws_environment_creds: + session_kwargs = { + "aws_access_key_id": self._aws_environment_creds["aws_access_key_id"], + "aws_secret_access_key": self._aws_environment_creds["aws_secret_access_key"], + } + if self._aws_environment_creds.get("aws_session_token"): + session_kwargs["aws_session_token"] = self._aws_environment_creds["aws_session_token"] + if region: + session_kwargs["region_name"] = region + elif os.environ.get("AWS_ACCESS_KEY_ID") or os.environ.get("AWS_SECRET_ACCESS_KEY"): + raise ValueError("Refusing live AWS credential environment fallback") + session = boto3.Session(**session_kwargs) + if not self._aws_environment_creds: + self._bind_aws_workload_token_sources(session) + if not self._aws_environment_creds and self._aws_profile_uses_credential_process(session): + if not self._aws_static_creds: + raise ValueError("AWS credential_process is incompatible with request isolation") + self._aws_active_creds = dict(self._aws_static_creds) + self._aws_imds_fell_back = True + get_logger().warning( + "AWS_USE_IMDS: credential_process is incompatible with request isolation; " + "using static credentials" + ) + return False + if not region: + try: + region = session.region_name + except Exception as e: + get_logger().warning(f"AWS_USE_IMDS: failed to resolve region via boto3: {e}") + creds = session.get_credentials() + if creds: + frozen_credentials = creds.get_frozen_credentials() + self._validate_aws_credential_chain_environment() + self._aws_boto3_creds = creds + self._aws_active_creds = self._aws_request_params_from_frozen(frozen_credentials, region) + self._aws_imds_mode = True + get_logger().info("Using ambient AWS credentials from IMDS/task-role/IRSA") + else: + get_logger().warning( + "AWS_USE_IMDS is set but boto3 found no credentials; falling through to static keys" + ) + except (botocore.exceptions.BotoCoreError, botocore.exceptions.ClientError, OSError): + get_logger().exception( + "AWS_USE_IMDS: failed to resolve credentials via boto3; falling through to static keys" + ) + + if not region: + get_logger().warning("AWS_USE_IMDS: could not determine AWS region; set AWS_REGION_NAME explicitly") + if not self._aws_imds_mode and self._aws_static_creds: + self._aws_active_creds = dict(self._aws_static_creds) + self._aws_imds_fell_back = True + get_logger().info("AWS_USE_IMDS: IMDS resolution failed; using static credentials") + return self._aws_imds_mode + + def _bind_aws_workload_token_sources(self, session) -> None: + """Bind token selectors without freezing deployment-owned token rotation.""" + from botocore.credentials import AssumeRoleWithWebIdentityProvider, ContainerProvider + + environment = { + key: value for key, value in self._aws_credential_chain_environment.items() if value is not None + } + resolver = session._session.get_component("credential_provider") + # Adapt only this session's botocore providers. Its EcsContainer role + # sourcer shares the container provider; nested web identity profiles + # already disable environment selection in botocore 1.43.x. + for provider in resolver.providers: + if isinstance(provider, ContainerProvider): + provider._environ = dict(environment) + elif isinstance(provider, AssumeRoleWithWebIdentityProvider) and not provider._disable_env_vars: + values = { + key: environment[variable] + for key, variable in provider._CONFIG_TO_ENV_VAR.items() if variable in environment + } + + def get_env_config(key, values=values): + return values.get(key) + + provider._get_env_config = get_env_config @staticmethod - def _litellm_otel_callback_enabled() -> bool: - """True when litellm's built-in OpenTelemetry callback is registered.""" - return any( - "otel" in (getattr(litellm, name, None) or []) - for name in ("callbacks", "success_callback", "failure_callback", "service_callback") - ) + def _aws_profile_uses_credential_process(session) -> bool: + """Return whether the selected boto3 profile chain executes a credential process.""" + profile_name = getattr(session, "profile_name", None) + botocore_session = getattr(session, "_session", None) + full_config = getattr(botocore_session, "full_config", None) + if not isinstance(profile_name, str) or not isinstance(full_config, dict): + return False + + profiles = full_config.get("profiles", {}) + visited_profiles = set() + while profile_name and profile_name not in visited_profiles: + visited_profiles.add(profile_name) + profile = profiles.get(profile_name, {}) + if not isinstance(profile, dict): + return False + if profile.get("credential_process"): + return True + profile_name = profile.get("source_profile") + return False @staticmethod - def _write_frozen_aws_creds_to_env(frozen) -> None: - """Write a botocore FrozenCredentials snapshot into os.environ for litellm/Bedrock.""" - os.environ["AWS_ACCESS_KEY_ID"] = frozen.access_key - os.environ["AWS_SECRET_ACCESS_KEY"] = frozen.secret_key - if frozen.token: - os.environ["AWS_SESSION_TOKEN"] = frozen.token - elif "AWS_SESSION_TOKEN" in os.environ: - del os.environ["AWS_SESSION_TOKEN"] + def _aws_request_params_from_frozen(frozen, region) -> dict: + """Convert a botocore credential snapshot to LiteLLM request parameters.""" + params = { + "aws_access_key_id": frozen.access_key, + "aws_secret_access_key": frozen.secret_key, + "aws_session_token": frozen.token or "", + } + if region: + params["aws_region_name"] = region + return params def _refresh_aws_imds_credentials(self) -> bool: """Refresh ambient AWS credentials from boto3 provider chain. Called before each Bedrock call to avoid serving stale credentials from long-lived processes (EC2 roles rotate every ~6h). - Uses the credentials object stored during __init__ rather than creating a new boto3.Session, - which would read the already-set AWS_* env vars and return stale values. + Uses the credentials object stored during initial ambient resolution rather than creating a new boto3.Session. Returns True on success, False on failure (caller should trigger static fallback).""" import botocore.exceptions @@ -416,24 +2671,609 @@ def _refresh_aws_imds_credentials(self) -> bool: if self._aws_boto3_creds is None: get_logger().warning("IMDS credential refresh: no boto3 credentials object stored") return False - self._write_frozen_aws_creds_to_env(self._aws_boto3_creds.get_frozen_credentials()) + self._validate_aws_credential_chain_environment() + region = self._aws_active_creds.get("aws_region_name") + frozen_credentials = self._aws_boto3_creds.get_frozen_credentials() + self._validate_aws_credential_chain_environment() + self._aws_active_creds = self._aws_request_params_from_frozen(frozen_credentials, region) return True - except (botocore.exceptions.BotoCoreError, botocore.exceptions.ClientError): + except (botocore.exceptions.BotoCoreError, botocore.exceptions.ClientError, OSError): # ClientError (STS/AssumeRole failures) is not a BotoCoreError subclass. get_logger().exception("IMDS credential refresh failed") return False def _activate_static_aws_fallback(self): - """Swap process env to static credentials for Bedrock fallback after IMDS failure.""" - os.environ["AWS_ACCESS_KEY_ID"] = self._aws_static_creds["AWS_ACCESS_KEY_ID"] - os.environ["AWS_SECRET_ACCESS_KEY"] = self._aws_static_creds["AWS_SECRET_ACCESS_KEY"] - os.environ["AWS_REGION_NAME"] = self._aws_static_creds["AWS_REGION_NAME"] - if "AWS_SESSION_TOKEN" in self._aws_static_creds: - os.environ["AWS_SESSION_TOKEN"] = self._aws_static_creds["AWS_SESSION_TOKEN"] - elif "AWS_SESSION_TOKEN" in os.environ: - del os.environ["AWS_SESSION_TOKEN"] + """Select static request credentials for an AWS provider fallback after IMDS failure.""" + self._aws_active_creds = dict(self._aws_static_creds) self._aws_imds_fell_back = True - get_logger().warning("Bedrock call failed with ambient (IMDS) credentials; retrying with static credentials") + get_logger().warning("AWS provider call failed with ambient credentials; retrying with static credentials") + + def _validate_aws_credential_chain_environment(self) -> None: + """Reject credential-chain selectors changed after this handler was initialized.""" + if any( + os.environ.get(variable) != value + for variable, value in self._aws_credential_chain_environment.items() + ): + raise ValueError("Refusing changed AWS credential-chain environment") + if ( + self._aws_credential_chain_files + and self._snapshot_aws_credential_chain_files() != self._aws_credential_chain_files + ): + raise ValueError("Refusing changed AWS credential-chain file") + + def _validate_aws_request_endpoint_environment(self) -> None: + """Reject AWS request endpoints changed after this handler was initialized.""" + if any( + os.environ.get(variable) != self._aws_credential_chain_environment.get(variable) + for variable in AWS_REQUEST_ENDPOINT_ENV_VARS + ): + raise ValueError("Refusing changed AWS request endpoint environment") + + @staticmethod + def _original_ec2_credential_file_path() -> str | None: + """Return the path that botocore's OriginalEC2Provider would read.""" + if original_ec2_credentials := os.environ.get("AWS_CREDENTIAL_FILE"): + return os.path.abspath(os.path.expanduser(original_ec2_credentials)) + return None + + @classmethod + def _aws_credential_chain_file_paths(cls) -> tuple[str, ...]: + """Return credential files that boto3 may read for this handler.""" + home = os.path.expanduser("~") + paths = [] + for variable, default_path in ( + ("AWS_SHARED_CREDENTIALS_FILE", os.path.join(home, ".aws", "credentials")), + ("AWS_CONFIG_FILE", os.path.join(home, ".aws", "config")), + ): + configured_path = os.environ.get(variable) + if configured_path is None: + paths.append(default_path) + elif configured_path: + paths.append(configured_path) + boto_config = os.environ.get("BOTO_CONFIG") + if boto_config is None: + paths.extend(("/etc/boto.cfg", os.path.join(home, ".boto"))) + elif boto_config: + paths.append(boto_config) + normalized_paths = [ + os.path.abspath(os.path.expanduser(os.path.expandvars(path))) + for path in paths + ] + if original_ec2_credentials := cls._original_ec2_credential_file_path(): + normalized_paths.append(original_ec2_credentials) + return tuple(dict.fromkeys(normalized_paths)) + + @staticmethod + def _fingerprint_aws_credential_chain_file(path: str) -> tuple: + """Return a stable fingerprint without retaining credential file contents.""" + try: + if not stat.S_ISREG(os.stat(path).st_mode): + return "nonfile", + with open(path, "rb") as credential_file: + return "file", hashlib.file_digest(credential_file, "sha256").hexdigest() + except FileNotFoundError: + return "missing", + except OSError as error: + return "error", type(error).__name__, error.errno + + @classmethod + def _snapshot_aws_credential_chain_files(cls) -> dict: + """Capture fingerprints for boto3 credential-chain files.""" + fingerprints = { + path: cls._fingerprint_aws_credential_chain_file(path) + for path in cls._aws_credential_chain_file_paths() + } + original_ec2_credentials = cls._original_ec2_credential_file_path() + if original_ec2_credentials and fingerprints.get(original_ec2_credentials) == ("nonfile",): + raise ValueError("AWS_CREDENTIAL_FILE must reference a regular file") + return fingerprints + + @contextlib.asynccontextmanager + async def _snapshot_aws_request_credentials(self, enabled): + """Refresh synchronously and serialize this handler's AWS call and static fallback.""" + if not enabled: + yield dict(self._aws_active_creds), False + return + async with self._aws_bedrock_lock: + if not self._aws_imds_fell_back: + self._validate_aws_credential_chain_environment() + if self._aws_imds_mode and not self._refresh_aws_imds_credentials() and self._aws_static_creds: + self._activate_static_aws_fallback() + can_fallback = self._aws_imds_mode and not self._aws_imds_fell_back and bool(self._aws_static_creds) + yield dict(self._aws_active_creds), can_fallback + + def _should_use_aws_imds(self, provider: str | None) -> bool: + """Return whether this request needs SigV4 credentials from the ambient AWS chain.""" + if not getattr(self, "_aws_use_imds", False) or provider not in AWS_REQUEST_PROVIDERS: + return False + if provider not in ("bedrock", "bedrock_mantle"): + return True + provider_params = getattr(self, "_provider_request_params", {}).get(provider, {}) + provider_environment_api_keys = getattr(self, "_provider_environment_api_keys", {}) + return not (provider_params.get("api_key") or provider_environment_api_keys.get(provider)) + + def _resolve_request_provider(self, model: str) -> str | None: + """Resolve the LiteLLM provider after PR-Agent has normalized the model name.""" + if not isinstance(model, str) or not model: + return None + provider_cache = getattr(self, "_request_provider_cache", None) + if provider_cache is None: + provider_cache = self._request_provider_cache = {} + if model in provider_cache: + return provider_cache[model] + transport_provider = None + if "/" in model: + transport_provider = model.split("/", 1)[0] + provider = PROVIDER_SETTING_ALIASES.get(transport_provider, transport_provider) + provider_params = getattr(self, "_provider_request_params", {}) + if ( + provider in provider_params + or provider in PROVIDER_API_KEY_ENV_VARS + or provider in AWS_REQUEST_PROVIDERS + or provider in OPENAI_COMPATIBLE_REQUEST_PROVIDERS + or _is_openai_compatible_request_provider(provider) + or provider in getattr(litellm, "provider_list", ()) + or provider in ("azure", "databricks", "openai", "vertex_ai") + ): + resolved_provider = provider + else: + resolved_provider = None + else: + try: + _, transport_provider, _, _ = litellm.get_llm_provider(model=model) + resolved_provider = PROVIDER_SETTING_ALIASES.get(transport_provider, transport_provider) + except litellm.BadRequestError: + if model.startswith("claude"): + resolved_provider = "anthropic" + elif model.startswith("command"): + resolved_provider = "cohere_chat" + else: + resolved_provider = "openai" + transport_provider = resolved_provider + provider_cache[model] = resolved_provider + transport_provider_cache = getattr(self, "_request_transport_provider_cache", None) + if transport_provider_cache is None: + transport_provider_cache = self._request_transport_provider_cache = {} + transport_provider_cache[model] = transport_provider + return resolved_provider + + def _resolve_request_transport_provider(self, model: str) -> str | None: + """Resolve the unaliased provider LiteLLM uses to select a transport.""" + self._resolve_request_provider(model) + return getattr(self, "_request_transport_provider_cache", {}).get(model) + + def _route_model(self, model: str, deployment_id: str | None) -> str: + """Apply provider routing shared by regular calls and health probes.""" + if model.startswith("azure_text/") and deployment_id: + return f"azure_text/{deployment_id}" + if self.azure: + if "/" not in model: + return "azure/" + model + provider, model_name = model.split("/", 1) + if provider == "azure_text" or provider == "openai" or PROVIDER_SETTING_ALIASES.get(provider) == "openai": + azure_provider = "azure_text" if provider in ("text-completion-openai", "azure_text") else "azure" + if azure_provider == "azure_text" and deployment_id: + model_name = deployment_id + return f"{azure_provider}/{model_name}" + return model + + def _route_model_for_request( + self, + model: str, + custom_llm_provider: str, + deployment_id: str | None, + ) -> str: + """Apply automatic routing while preserving explicit custom-provider model IDs.""" + if not custom_llm_provider: + model = self._route_model(model, deployment_id) + elif deployment_id and (custom_llm_provider == "azure_text" or model.startswith("azure_text/")): + model = f"azure_text/{deployment_id}" + return normalize_litellm_model(model, custom_llm_provider) + + @staticmethod + def _canonical_openrouter_model(model: str, provider: str | None) -> str | None: + """Return an OpenRouter-prefixed model for request-control matching.""" + if provider != "openrouter" or not isinstance(model, str): + return None + return model if model.startswith("openrouter/") else f"openrouter/{model}" + + @staticmethod + def _is_gpt6_astra_model(model: str) -> bool: + """Recognize native Astra models without changing gateway model IDs.""" + while model.startswith(("openai/", "azure/")): + model = model.removeprefix("openai/").removeprefix("azure/") + return model.removesuffix("_thinking") == "gpt-6-astra" + + @staticmethod + def _is_gpt5_model(model: str) -> bool: + """Return whether a routed model belongs to the GPT-5 family.""" + model_base = model.removeprefix("openrouter/") + while model_base.startswith(("openai/", "azure/")): + model_base = model_base.removeprefix("openai/").removeprefix("azure/") + return model_base.startswith("gpt-5") + + def _normalize_gpt5_model_for_request(self, model: str, user_model: str, custom_llm_provider: str) -> str: + """Normalize GPT-5/Astra suffixes and prefixes before request parameters are selected.""" + model_base = model + while model_base.startswith(("openai/", "azure/")): + model_base = model_base.removeprefix("openai/").removeprefix("azure/") + if not model_base.startswith("gpt-5") and model_base.removesuffix("_thinking") != "gpt-6-astra": + return model + if custom_llm_provider: + return model.replace("_thinking", "") + if self.azure or user_model.startswith("azure/"): + provider_prefix = "azure/" + else: + provider_prefix = "openai/" + return provider_prefix + model_base.replace("_thinking", "") + + def _uses_captured_azure_companion_auth(self, provider): + if provider not in ("azure", "azure_ai") or not getattr(self, "_azure_companion_auth", False): + return False + if provider == "azure_ai": + # This bridge covers initially tokenless raw companion selection, + # not another Azure key or AD-token choice in the common validator. + snapshot = self._raw_guard_auth_snapshot + return not (snapshot["azure_key"] or snapshot["azure_ad_token"]) + return True + + def _get_provider_request_params( + self, + model: str, + azure_ad_token=None, + provider=None, + transport_provider=None, + transport_model=None, + aws_request_credentials=None, + ) -> dict: + """Return only the credentials and routing parameters for this model's provider.""" + provider = provider or self._resolve_request_provider(model) + if provider in AWS_REQUEST_PROVIDERS: + self._validate_aws_request_endpoint_environment() + if transport_provider is None: + transport_provider = self._resolve_request_transport_provider(model) or provider + provider_params = getattr(self, "_provider_request_params", {}) + provider_environment_api_keys = getattr(self, "_provider_environment_api_keys", {}) + openai_api_base_is_azure = getattr(self, "_openai_api_base_is_azure", getattr(self, "azure", False)) + if provider in OPENAI_COMPATIBLE_REQUEST_PROVIDERS: + openai_params = {} if openai_api_base_is_azure else provider_params.get("openai", {}) + params = { + key: openai_params[key] + for key in ("api_key", "api_base") + if openai_params.get(key) + } + native_api_key = provider_environment_api_keys.get(provider) + if provider == "openai_like" and provider_params.get(provider, {}).get("api_base"): + params["api_base"] = provider_params[provider]["api_base"] + if params["api_base"] != openai_params.get("api_base"): + params.pop("api_key", None) + if provider == "openai_like" and native_api_key: + params["api_key"] = native_api_key + elif "api_key" not in params and native_api_key: + params["api_key"] = native_api_key + if provider == "openai_like" and "api_key" not in params and ( + openai_params.get("api_key") + or getattr(litellm, "api_key", None) + or getattr(litellm, "openai_key", None) + or _has_provider_api_key_global(provider) + or provider_environment_api_keys.get("openai") + or _has_live_provider_api_key_environment(provider) + or os.environ.get("OPENAI_API_KEY") + ): + params["api_key"] = DUMMY_LITELLM_API_KEY + if provider == "custom_openai" and "api_key" not in params: + params["api_key"] = DUMMY_LITELLM_API_KEY + if provider == "custom_openai" and "api_base" not in params: + params["api_base"] = OPENAI_DEFAULT_API_BASE + request_headers = _request_local_openai_headers(transport_provider, model=transport_model or model) + if request_headers is not None: + params["headers"] = request_headers + return self._finalize_provider_request_params(provider, params) + if provider is None: + # OPENAI.API_BASE also configures gateways such as MOSAICO whose model + # names can retain another provider prefix. Forward only the matching + # request-local key to that explicitly configured endpoint. + openai_params = ( + {} + if openai_api_base_is_azure + else getattr(self, "_provider_request_params", {}).get("openai", {}) + ) + api_base = openai_params.get("api_base") + params = {"api_base": api_base} if api_base else {} + request_api_key = openai_params.get("api_key") or provider_environment_api_keys.get("openai") + if api_base: + params["api_key"] = request_api_key or DUMMY_LITELLM_API_KEY + elif ( + getattr(litellm, "api_key", None) + or getattr(litellm, "openai_key", None) + or os.environ.get("OPENAI_API_KEY") + or getattr(openai, "api_key", None) + ): + params["api_key"] = DUMMY_LITELLM_API_KEY + return self._finalize_provider_request_params(provider, params) + if provider in MANAGED_AUTH_REQUEST_PROVIDERS: + params = dict(provider_params.get(provider, {})) + params["api_key"] = DUMMY_LITELLM_API_KEY + request_headers = _request_local_openai_headers(transport_provider, model=transport_model or model) + if request_headers is not None: + params["headers"] = request_headers + return self._finalize_provider_request_params(provider, params) + params = dict(provider_params.get(provider, {})) + if provider == "gdc": + api_base = params.get("api_base") + if not api_base: + raise ValueError("GDC API base was not resolved for this request; set GDC_API_BASE") + if "/v1/projects/" not in api_base: + from litellm.llms.gdc.chat.transformation import GDCGeminiConfig + + routing = { + name: provider_params.get("vertex_ai", {}).get(name) + for name in ("vertex_project", "vertex_location") + } + if not all(routing.values()): + raise ValueError("GDC host routing requires captured Vertex project and location") + # Native dispatch places routing kwargs below process globals. + # Resolve the full URL first, using the native highest-priority + # inputs, so authentication and quota headers share this route. + params["api_base"] = GDCGeminiConfig().get_complete_url( + api_base=api_base, api_key=params.get("api_key"), model=transport_model or model, + optional_params={}, litellm_params=routing, + ) + if provider == "vertex_ai": + if self._vertex_credentials_error: + raise ValueError( + f"Unable to snapshot explicit Vertex credentials: {self._vertex_credentials_error}" + ) + if self._vertex_credentials is None: + captured_path = (self._vertex_default_adc or {}).get("explicit_path") + live_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") + if os.environ.get("VERTEXAI_CREDENTIALS") or (live_path and live_path != captured_path): + raise ValueError("Refusing live Vertex credential environment fallback") + if provider == "openrouter" and not params and not provider_environment_api_keys.get("openrouter"): + openai_params = {} if openai_api_base_is_azure else provider_params.get("openai", {}) + gateway_api_base = openai_params.get("api_base") + gateway_api_key = openai_params.get("api_key") or provider_environment_api_keys.get("openai") + if gateway_api_base and gateway_api_key: + params.update(api_base=gateway_api_base, api_key=gateway_api_key) + if provider == "openrouter" and "api_base" not in params and ( + params.get("api_key") + or provider_environment_api_keys.get("openrouter") + or _has_live_provider_api_key_environment("openrouter") + ): + params["api_base"] = "https://openrouter.ai/api/v1" + volcengine_responses = provider == "volcengine" and _uses_openai_responses_transport( + transport_model or model, transport_provider, + ) + if volcengine_responses: + # Native provider resolution prefers VOLCENGINE_API_KEY; only the + # Responses transport falls back to ARK_API_KEY afterward. + if "api_key" not in params: + api_key = provider_environment_api_keys.get(provider) or getattr(self, "_volcengine_ark_api_key", None) + if api_key: + params["api_key"] = api_key + if ( + _is_openai_compatible_request_provider(provider) + and not params + and not provider_environment_api_keys.get(provider) + ): + # Use OPENAI.API_BASE only when this provider has no native request + # identity. This prevents sending provider credentials to a gateway. + api_base = None if openai_api_base_is_azure else provider_params.get("openai", {}).get("api_base") + if api_base: + params["api_base"] = api_base + params["api_key"] = ( + provider_params.get("openai", {}).get("api_key") + or provider_environment_api_keys.get("openai") + or DUMMY_LITELLM_API_KEY + ) + if volcengine_responses and "api_key" not in params and os.environ.get("ARK_API_KEY"): + raise ValueError("Refusing Volcengine ARK API key added after handler initialization") + if provider in ("watsonx", "watsonx_text") and "api_key" not in params and ( + params.get("token") or params.get("zen_api_key") + ): + # LiteLLM resolves a Watsonx API key before applying explicit token + # or Zen authentication. Block that ambient fallback request-locally. + params["api_key"] = DUMMY_LITELLM_API_KEY + if "api_key" not in params: + api_key = provider_environment_api_keys.get(provider) + if api_key: + params["api_key"] = api_key + if provider == "snowflake": + # Require captured identity before the generic placeholder guard; + # native Snowflake cannot authenticate a keyless request. + if not params.get("api_key"): + raise ValueError("Snowflake JWT was not resolved for this request; set SNOWFLAKE_JWT") + if not params.get("api_base"): + account_id = getattr(self, "_snowflake_account_id", "") + if not account_id: + raise ValueError("Snowflake account was not resolved for this request; set SNOWFLAKE_ACCOUNT_ID") + params["account_id"] = account_id + if provider in ("hosted_vllm", "lm_studio") and params.get("api_base") and "api_key" not in params: + params["api_key"] = DUMMY_LITELLM_API_KEY + # Unlike Mantle, classic Bedrock never consumes litellm.api_key. + if provider == "bedrock" and "api_key" not in params and _has_live_provider_api_key_environment(provider): + raise ValueError("Refusing process-wide Bedrock bearer token fallback") + if provider in ("sagemaker_chat", "sagemaker_nova") and os.environ.get("AWS_BEARER_TOKEN_BEDROCK"): + # LiteLLM 1.100.0's SageMaker signer ignores its api_key argument and + # otherwise reads this Bedrock-only token directly from the environment. + raise ValueError("Refusing Bedrock bearer token fallback for SageMaker") + if provider == "azure" and getattr(self, "_azure_ad", False): + if azure_ad_token is None: + raise ValueError("Azure AD token was not resolved for this request") + params.pop("api_key", None) + params["azure_ad_token"] = azure_ad_token + if provider == "azure" and "azure_ad_token" not in params and any( + os.environ.get(environment_variable) for environment_variable in AZURE_AD_TOKEN_ENV_VARS + ): + raise ValueError("Refusing Azure AD token added after handler initialization") + if provider == "anthropic" and "api_key" not in params: + # LiteLLM has no completion parameter for ANTHROPIC_AUTH_TOKEN. The + # placeholder blocks ambient API keys until the request-local bridge + # restores Anthropic's keyless state and supplies the bearer token. + params["api_key"] = DUMMY_LITELLM_API_KEY + requires_api_key_guard = provider == "bedrock_mantle" or ( + provider not in AWS_REQUEST_PROVIDERS + and provider != "vertex_ai" + and (_uses_provider_api_key(provider) or provider in getattr(litellm, "provider_list", ())) + ) + companion_auth = ( + self._uses_captured_azure_companion_auth(provider) + and not _is_cloudflare_gateway(params.get("api_base")) + ) + # Token- or companion-based Azure requests must block even a key that appears only + # after native dispatch starts, not just an already visible fallback. + if "api_key" not in params and requires_api_key_guard and ( + (provider == "azure" and params.get("azure_ad_token")) + or companion_auth + or getattr(litellm, "api_key", None) + or _has_provider_api_key_global(provider) + or ( + _is_openai_compatible_request_provider(provider) + and (getattr(litellm, "openai_key", None) or provider_environment_api_keys.get("openai")) + ) + or _has_live_provider_api_key_environment(provider) + ): + params["api_key"] = DUMMY_LITELLM_API_KEY + params["_raw_api_key_guard"] = True + if provider == "azure" and str(params.get("azure_ad_token") or "").startswith("oidc/"): + params["_azure_oidc_guard"] = True + elif provider == "azure" and (params.get("azure_ad_token") or companion_auth): + params["_azure_ad_guard"] = True + if provider == "openai" and "api_key" not in params: + # Do not trust process-wide LiteLLM globals here: another request or + # embedding may have populated them with a different tenant's key. + params["api_key"] = DUMMY_LITELLM_API_KEY + if provider == "openai" and "api_base" not in params: + params["api_base"] = OPENAI_DEFAULT_API_BASE + if provider == "gemini": + # Freeze absence before LiteLLM moves dispatch to an executor, where + # a later key/base would otherwise replace the native default route. + params.setdefault("api_key", DUMMY_LITELLM_API_KEY) + if "api_base" not in params: + from litellm.llms.vertex_ai.common_utils import _get_gemini_url + + native_model = (transport_model or model).removeprefix("gemini/") + url, endpoint = _get_gemini_url(mode="chat", model=native_model, stream=False) + suffix = f"/models/{native_model}:{endpoint}" + if not url.endswith(suffix) or len(url) == len(suffix): + raise RuntimeError("LiteLLM's Gemini default URL is incompatible with request isolation") + params["api_base"] = url[:-len(suffix)] + if provider in AWS_REQUEST_PROVIDERS: + model_region = ( + _get_bedrock_model_region(transport_model or model, getattr(self, "_bedrock_model_id", None)) + if provider == "bedrock" else None + ) + if model_region: + # Native model regions precede environment/settings regions; + # capturing credentials must not promote their region above it. + params["aws_region_name"] = model_region + uses_bedrock_bearer = ( + provider in ("bedrock", "bedrock_mantle") + and params.get("api_key") not in (None, DUMMY_LITELLM_API_KEY) + ) + if not uses_bedrock_bearer: + if any(os.environ.get(variable) for variable in LITELLM_AWS_CREDENTIAL_SELECTOR_ENV_VARS): + # LiteLLM 1.100.0 resolves these selectors ahead of explicit + # request credentials, which would replace the isolated keys. + raise ValueError(f"Refusing ambient LiteLLM AWS credential selector for provider {provider}") + aws_request_credentials = dict(aws_request_credentials or {}) + if (provider == "bedrock_mantle" or model_region) and params.get("aws_region_name"): + aws_request_credentials["aws_region_name"] = params["aws_region_name"] + if not ( + aws_request_credentials.get("aws_access_key_id") + and aws_request_credentials.get("aws_secret_access_key") + ): + if getattr(self, "_aws_environment_credentials_incomplete", False): + raise ValueError("AWS environment credentials are incomplete") + raise ValueError("AWS credentials were not resolved for this request") + if not aws_request_credentials.get("aws_region_name"): + raise ValueError( + "AWS region was not resolved for this request; set AWS_REGION_NAME, " + "aws.AWS_REGION_NAME, AWS_REGION, or AWS_DEFAULT_REGION" + ) + params.update(aws_request_credentials) + request_headers = _request_local_openai_headers( + transport_provider, + params.get("organization"), + transport_model or model, + ) + if request_headers is not None: + params["headers"] = request_headers + return self._finalize_provider_request_params(provider, params) + + def _finalize_provider_request_params(self, provider: str | None, params: dict) -> dict: + """Merge request-local headers and reject LiteLLM's process-wide header fallback.""" + params = _guard_request_routing_globals(provider, params) + watsonx_token = params.pop("token", None) if provider in ("watsonx", "watsonx_text") else None + request_headers = dict(getattr(self, "_request_headers", {})) + for header, value in (params.get("headers") or {}).items(): + matching_headers = [name for name in request_headers if name.lower() == header.lower()] + if matching_headers and isinstance(value, openai.Omit): + value = request_headers[matching_headers[-1]] + for matching_header in matching_headers: + del request_headers[matching_header] + request_headers[header] = value + authorization_headers = [header for header in request_headers if header.lower() == "authorization"] + has_authorization = bool(authorization_headers) + if provider in ("watsonx", "watsonx_text") and has_authorization: + authorization = request_headers[authorization_headers[-1]] + for header in authorization_headers: + del request_headers[header] + request_headers["Authorization"] = authorization + if watsonx_token and not has_authorization: + request_headers["Authorization"] = f"Bearer {watsonx_token}" + has_authorization = True + if provider in ("watsonx", "watsonx_text") and has_authorization: + params.pop("zen_api_key", None) + params["api_key"] = DUMMY_LITELLM_API_KEY + if request_headers: + params["headers"] = request_headers + elif getattr(litellm, "headers", None): + raise ValueError(f"Refusing process-wide LiteLLM headers fallback for provider {provider or 'unknown'}") + return params + + def _requires_streaming(self, model: str) -> bool: + """Return whether this model requires streaming after OpenAI/Azure routing.""" + def normalize(candidate: str) -> str: + while candidate.startswith(("azure/", "openai/")): + candidate = candidate.removeprefix("azure/").removeprefix("openai/") + return candidate + + normalized_model = normalize(model) + return any(normalize(candidate) == normalized_model for candidate in self.streaming_required_models) + + def _force_streaming_for_request(self, custom_llm_provider, api_base) -> bool: + """Return whether an OpenAI-compatible endpoint requires streaming.""" + custom_llm_provider = str(custom_llm_provider or "").strip().lower() + api_base = api_base.strip().lower() if isinstance(api_base, str) else "" + return ( + bool(custom_llm_provider) + and custom_llm_provider == self.force_streaming_provider + and bool(self.force_streaming_api_base_substrings) + and any(substring in api_base for substring in self.force_streaming_api_base_substrings) + ) + + async def _get_provider_request_params_async( + self, + model: str, + provider=None, + transport_provider=None, + transport_model=None, + aws_request_credentials=None, + ) -> dict: + """Resolve provider parameters without blocking the event loop on Azure AD refresh.""" + provider = provider or self._resolve_request_provider(model) + azure_ad_token = None + if getattr(self, "_azure_ad", False) and provider == "azure": + azure_ad_token = await asyncio.to_thread(_get_azure_ad_token, self._azure_ad_credential) + return self._get_provider_request_params( + model, + azure_ad_token=azure_ad_token, + provider=provider, + transport_provider=transport_provider, + transport_model=transport_model, + aws_request_credentials=aws_request_credentials, + ) def prepare_logs(self, response, system, user, resp, finish_reason): response_log = response.dict().copy() @@ -543,6 +3383,251 @@ def _clamp_grok_reasoning_effort(cls, model: str, reasoning_effort: str) -> str: return "xhigh" if "xhigh" in grok_levels else "high" return "low" + def _resolve_reasoning_effort(self, model: str, configured_effort) -> str: + """Validate and normalize a configured reasoning effort for this model.""" + try: + ReasoningEffort(configured_effort) + reasoning_effort = configured_effort + except (ValueError, TypeError): + reasoning_effort = ReasoningEffort.MEDIUM.value + if configured_effort is not None: + get_logger().warning( + f"Invalid reasoning_effort '{configured_effort}' in config. " + f"Using default '{reasoning_effort}'. Valid values: {[e.value for e in ReasoningEffort]}" + ) + + clamped_effort = self._clamp_grok_reasoning_effort(model, reasoning_effort) + if clamped_effort != reasoning_effort: + get_logger().info( + f"Grok model {model} does not support reasoning_effort='{reasoning_effort}'; " + f"using '{clamped_effort}' instead." + ) + return clamped_effort + + def _apply_openrouter_request_controls( + self, + model: str, + kwargs: dict, + inherited_reasoning_effort: str | None = None, + ) -> dict: + """Apply handler-local OpenRouter routing, reasoning, and output controls.""" + openrouter_settings = self._openrouter_controls + extra_body = kwargs.get("extra_body") or {} + + # Normalize operator-controlled config: Dynaconf/env overrides can + # arrive as strings (AUTO_CAST_FOR_DYNACONF is disabled), so coerce + # defensively instead of trusting the declared types. + provider_only = _as_list(openrouter_settings.get("provider_only", [])) + provider_order = _as_list(openrouter_settings.get("provider_order", [])) + if provider_only: + extra_body.setdefault("provider", {})["only"] = provider_only + elif provider_order: + provider = extra_body.setdefault("provider", {}) + provider["order"] = provider_order + provider["allow_fallbacks"] = _as_bool(openrouter_settings.get("allow_fallbacks", True), default=True) + + reasoning = {} + effective_reasoning_effort = str( + openrouter_settings.get("reasoning_effort", "") or "" + ).strip().lower() + reasoning_max_tokens = self._coerce_token_value(openrouter_settings.get("reasoning_max_tokens", 0)) + if effective_reasoning_effort: + try: + ReasoningEffort(effective_reasoning_effort) + except (TypeError, ValueError): + get_logger().warning( + f"Ignoring invalid openrouter.reasoning_effort '{effective_reasoning_effort}'. " + f"Valid values: {[effort.value for effort in ReasoningEffort]}." + ) + effective_reasoning_effort = "" + if not effective_reasoning_effort: + if reasoning_max_tokens > 0 and inherited_reasoning_effort: + if inherited_reasoning_effort == "none": + get_logger().warning( + f"Ignoring config.reasoning_effort='{inherited_reasoning_effort}' because " + "openrouter.reasoning_max_tokens takes precedence." + ) + else: + get_logger().info( + "Using openrouter.reasoning_max_tokens over" + f" config.reasoning_effort='{inherited_reasoning_effort}'." + ) + elif reasoning_max_tokens <= 0: + effective_reasoning_effort = inherited_reasoning_effort or "" + + if effective_reasoning_effort: + clamped_effort = self._clamp_grok_reasoning_effort(model, effective_reasoning_effort) + if clamped_effort != effective_reasoning_effort: + get_logger().info( + f"Grok model {model} does not support reasoning_effort=" + f"'{effective_reasoning_effort}'; using '{clamped_effort}' instead." + ) + effective_reasoning_effort = clamped_effort + + # Preserve explicit disablement; otherwise keep effort and max_tokens + # mutually exclusive by preferring the token budget. + if effective_reasoning_effort == "none": + if reasoning_max_tokens > 0: + get_logger().warning( + "Ignoring openrouter.reasoning_max_tokens because " + "openrouter.reasoning_effort='none' disables reasoning." + ) + reasoning["enabled"] = False + elif reasoning_max_tokens > 0: + if effective_reasoning_effort: + get_logger().warning( + f"Ignoring openrouter.reasoning_effort='{effective_reasoning_effort}' because " + "openrouter.reasoning_max_tokens takes precedence." + ) + reasoning["max_tokens"] = reasoning_max_tokens + elif effective_reasoning_effort: + # OpenRouter uses xhigh for the max alias; extra_body bypasses + # LiteLLM's OpenRouter parameter mapping. + reasoning["effort"] = "xhigh" if effective_reasoning_effort == "max" else effective_reasoning_effort + if reasoning: + get_logger().info(f"Adding OpenRouter reasoning {reasoning} to model {model}.") + extra_body["reasoning"] = reasoning + + if extra_body: + kwargs["extra_body"] = extra_body + + max_tokens = self._coerce_token_value(openrouter_settings.get("max_tokens", 0)) + output_limit_param = ( + "max_completion_tokens" if "max_completion_tokens" in kwargs else "max_tokens" + ) + if max_tokens > 0: + existing = self._coerce_token_value(kwargs.get(output_limit_param, 0)) + kwargs[output_limit_param] = min(existing, max_tokens) if existing > 0 else max_tokens + effective_max_tokens = self._coerce_token_value(kwargs.get(output_limit_param, 0)) + effective_reasoning_max_tokens = self._coerce_token_value(reasoning.get("max_tokens", 0)) + effective_reasoning_effort = reasoning.get("effort") + if ( + model.startswith("openrouter/anthropic/") + and 0 < effective_max_tokens + and ( + 0 < effective_reasoning_max_tokens >= effective_max_tokens + or (effective_reasoning_effort and effective_max_tokens <= 1024) + ) + ): + minimum_reasoning_tokens = effective_reasoning_max_tokens or 1024 + get_logger().warning( + f"OpenRouter Anthropic max_tokens ({effective_max_tokens}) must be greater than " + f"the reasoning budget ({minimum_reasoning_tokens}) to leave output headroom." + ) + return kwargs + + @staticmethod + def _coerce_token_value(value) -> int: + """Mirror the request's permissive integer coercion without conversion failures.""" + try: + return int(value) + except (TypeError, ValueError, OverflowError): + return 0 + + def _claude_thinking_mode(self, model: str) -> str | None: + """Return the thinking mode selected by request construction for this model.""" + adaptive_model = self._is_claude_adaptive_thinking_model(model) + if adaptive_model and self._claude_thinking_controls["enable_claude_adaptive_thinking"]: + return "adaptive" + if ( + model in self.claude_extended_thinking_models + and self._claude_thinking_controls["enable_claude_extended_thinking"] + ): + return "unsupported_extended" if adaptive_model else "extended" + return None + + def _get_claude_extended_thinking_limits(self) -> tuple[int, int]: + """Validate and return the snapshotted extended-thinking budget and output cap.""" + extended_thinking_budget_tokens = self._claude_thinking_controls["extended_thinking_budget_tokens"] + extended_thinking_max_output_tokens = self._claude_thinking_controls["extended_thinking_max_output_tokens"] + + if not isinstance(extended_thinking_budget_tokens, int) or extended_thinking_budget_tokens <= 0: + raise ValueError( + f"extended_thinking_budget_tokens must be a positive integer, " + f"got {extended_thinking_budget_tokens}" + ) + if not isinstance(extended_thinking_max_output_tokens, int) or extended_thinking_max_output_tokens <= 0: + raise ValueError( + f"extended_thinking_max_output_tokens must be a positive integer, " + f"got {extended_thinking_max_output_tokens}" + ) + if extended_thinking_max_output_tokens < extended_thinking_budget_tokens: + raise ValueError( + f"extended_thinking_max_output_tokens ({extended_thinking_max_output_tokens}) must be greater than " + f"or equal to extended_thinking_budget_tokens ({extended_thinking_budget_tokens})" + ) + return extended_thinking_budget_tokens, extended_thinking_max_output_tokens + + def _resolve_output_token_limit(self, model: str, openrouter_model: str | None) -> int: + """Return the final positive output cap selected by PR-Agent request controls.""" + output_tokens = self._coerce_token_value(get_settings().config.get("max_output_tokens", 0)) + if self._claude_thinking_mode(model) == "extended": + _, output_tokens = self._get_claude_extended_thinking_limits() + + if openrouter_model: + openrouter_output_tokens = self._coerce_token_value(self._openrouter_controls.get("max_tokens", 0)) + if openrouter_output_tokens > 0: + output_tokens = ( + min(output_tokens, openrouter_output_tokens) + if output_tokens > 0 + else openrouter_output_tokens + ) + return output_tokens if output_tokens > 0 else 0 + + def get_output_token_limit(self, model: str) -> int: + """Return the output cap that this handler will request for the supplied model.""" + custom_llm_provider = self._custom_llm_provider + configured_deployment_id = self.deployment_id + routed_model = self._route_model_for_request(model, custom_llm_provider, configured_deployment_id) + completion_model = self._normalize_gpt5_model_for_request(routed_model, model, custom_llm_provider) + request_provider = ( + PROVIDER_SETTING_ALIASES.get(custom_llm_provider, custom_llm_provider) + if custom_llm_provider + else self._resolve_request_provider(routed_model) + ) + openrouter_model = self._canonical_openrouter_model(completion_model, request_provider) + return self._resolve_output_token_limit(completion_model, openrouter_model) + + def get_output_token_reserve(self, model: str, default_output_tokens: int) -> int: + """Return completion headroom to reserve while fitting a request prompt.""" + output_tokens = self.get_output_token_limit(model) + if output_tokens > 0: + return output_tokens + + default_output_tokens = self._coerce_token_value(default_output_tokens) + custom_llm_provider = self._custom_llm_provider + routed_model = self._route_model_for_request(model, custom_llm_provider, self.deployment_id) + request_provider = ( + PROVIDER_SETTING_ALIASES.get(custom_llm_provider, custom_llm_provider) + if custom_llm_provider + else self._resolve_request_provider(routed_model) + ) + openrouter_model = self._canonical_openrouter_model( + routed_model, request_provider + ) + if not openrouter_model: + return default_output_tokens + + reasoning_effort = str( + self._openrouter_controls.get("reasoning_effort", "") or "" + ).strip().lower() + reasoning_effort = self._clamp_grok_reasoning_effort( + openrouter_model, reasoning_effort + ) + reasoning_tokens = self._coerce_token_value( + self._openrouter_controls.get("reasoning_max_tokens", 0) + ) + if reasoning_effort == "none" or reasoning_tokens <= 0: + return default_output_tokens + return default_output_tokens + reasoning_tokens + + @staticmethod + def normalize_request_prompts(model: str, system_prompt: str, user_prompt: str) -> tuple[str, str]: + """Return the prompt strings that request construction will send.""" + if 'claude' in model and not system_prompt: + system_prompt = "No system prompt provided" + return system_prompt, user_prompt + def _configure_claude_extended_thinking(self, model: str, kwargs: dict) -> dict: """ Configure Claude extended thinking parameters if applicable. @@ -554,16 +3639,9 @@ def _configure_claude_extended_thinking(self, model: str, kwargs: dict) -> dict: Returns: dict: Updated kwargs with extended thinking configuration """ - extended_thinking_budget_tokens = get_settings().config.get("extended_thinking_budget_tokens", 2048) - extended_thinking_max_output_tokens = get_settings().config.get("extended_thinking_max_output_tokens", 4096) - - # Validate extended thinking parameters - if not isinstance(extended_thinking_budget_tokens, int) or extended_thinking_budget_tokens <= 0: - raise ValueError(f"extended_thinking_budget_tokens must be a positive integer, got {extended_thinking_budget_tokens}") - if not isinstance(extended_thinking_max_output_tokens, int) or extended_thinking_max_output_tokens <= 0: - raise ValueError(f"extended_thinking_max_output_tokens must be a positive integer, got {extended_thinking_max_output_tokens}") - if extended_thinking_max_output_tokens < extended_thinking_budget_tokens: - raise ValueError(f"extended_thinking_max_output_tokens ({extended_thinking_max_output_tokens}) must be greater than or equal to extended_thinking_budget_tokens ({extended_thinking_budget_tokens})") + extended_thinking_budget_tokens, extended_thinking_max_output_tokens = ( + self._get_claude_extended_thinking_limits() + ) kwargs["thinking"] = { "type": "enabled", @@ -592,7 +3670,7 @@ def _is_claude_adaptive_thinking_model(model: str) -> bool: def _configure_claude_adaptive_thinking(self, model: str, kwargs: dict) -> dict: """Configure thinking for Claude models that reject token budgets.""" kwargs["thinking"] = {"type": "adaptive"} - effort = get_settings().config.reasoning_effort + effort = self._default_reasoning_effort if effort in ("low", "medium", "high", "xhigh", "max"): kwargs["output_config"] = {"effort": effort} get_logger().info( @@ -626,7 +3704,7 @@ def capture_logs(message): if extra.get("pr_url") is not None: log_entry.update({"pr_url": extra["pr_url"]}) - # Append the log entry to the captured_logs list + # Append the captured request context. captured_extra.append(log_entry) # Adding the custom sink to Loguru @@ -747,83 +3825,111 @@ def _resolve_cache_control_injection_points(): raise ValueError("LITELLM.CACHE_CONTROL_INJECTION_POINTS must be a JSON/TOML array") return cache_control_injection_points + async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2, img_path: str = None): + configured_deployment_id = self.deployment_id + return await self._chat_completion_with_retry( + model, + system, + user, + temperature, + img_path, + configured_deployment_id=configured_deployment_id, + ) + @retry( retry=retry_if_exception(_should_retry_same_model), stop=stop_after_attempt(MODEL_RETRIES), reraise=True, # surface the provider's error; RetryError hides the reason ) - async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2, img_path: str = None): - # Serialize env-var mutation + Bedrock call for IMDS mode to prevent concurrent - # requests from interleaving os.environ credentials during asyncio.gather usage. + async def _chat_completion_with_retry( + self, + model: str, + system: str, + user: str, + temperature: float = 0.2, + img_path: str = None, + *, + configured_deployment_id: str | None, + ): # Validate config-derived kwargs before the try/except below, so a malformed value raises a # ValueError config error instead of being wrapped as openai.APIError and retried. cache_control_injection_points = self._resolve_cache_control_injection_points() client_retries = _configured_client_retries() - _bedrock_imds = self._aws_imds_mode and any( - provider in model for provider in ("bedrock/", "bedrock_mantle/") + custom_llm_provider = self._custom_llm_provider + user_model = model + routed_model = self._route_model_for_request(user_model, custom_llm_provider, configured_deployment_id) + completion_model = self._normalize_gpt5_model_for_request(routed_model, user_model, custom_llm_provider) + request_provider = ( + PROVIDER_SETTING_ALIASES.get(custom_llm_provider, custom_llm_provider) + if custom_llm_provider + else self._resolve_request_provider(routed_model) ) - async with (self._aws_bedrock_lock if _bedrock_imds else contextlib.nullcontext()): - if _bedrock_imds and not self._aws_imds_fell_back: - if not self._refresh_aws_imds_credentials() and self._aws_static_creds: - self._activate_static_aws_fallback() - self._aws_imds_fell_back = True + deployment_id = ( + configured_deployment_id + if request_provider == "azure" and not routed_model.startswith("azure_text/") + else None + ) + if img_path: + try: + # Finish external image I/O before validating mutable credential fallbacks. + r = await asyncio.to_thread( + requests.head, + img_path, + allow_redirects=True, + timeout=_IMAGE_HEAD_TIMEOUT_SECONDS, + ) + if r.status_code == 404: + error_msg = "The image link is not [alive](img_path).\nPlease repost the original image as a comment, and send the question again with 'quote reply' (see [instructions](https://docs.pr-agent.ai/tools/ask/#ask-on-images))." + get_logger().error(error_msg) + return f"{error_msg}", "error" + except Exception as e: + get_logger().error(f"Error fetching image: {img_path}", e) + return f"Error fetching image: {img_path}", "error" + + _aws_imds = self._should_use_aws_imds(request_provider) + async with self._snapshot_aws_request_credentials(_aws_imds) as ( + aws_request_credentials, + aws_can_fallback, + ): + # Resolve credentials before the retry-wrapped inference block. In + # particular, Azure AD refresh failures are authentication errors and + # should not be converted to retryable OpenAI API errors below. + provider_request_params = await self._get_provider_request_params_async( + routed_model, + provider=request_provider, + transport_provider=custom_llm_provider or None, + transport_model=deployment_id or completion_model, + aws_request_credentials=aws_request_credentials, + ) try: resp, finish_reason = None, None - deployment_id = self.deployment_id - # Capture the original model string so an explicit provider prefix in the - # user config (e.g. "azure/gpt-5...") can be preserved when the GPT-5 branch - # rebuilds the routed model name below. - user_model = model - # Capture the provider prefix before any rewriting below. Databricks auth/endpoint - # selection keys off this; rewriting (e.g. 'azure/' + model when Azure is enabled in - # a multi-provider config) would otherwise hide the 'databricks/' prefix and bypass - # the guards that keep Databricks on its own DATABRICKS_API_KEY/DATABRICKS_API_BASE. - is_databricks = model.startswith("databricks/") - # OpenRouter models keep their "openrouter/" prefix: __init__ already - # routed them to the OpenRouter api_key/api_base, and rewriting to - # "azure/openrouter/..." would both misroute the request and skip the - # OpenRouter controls block below (guarded by the same prefix). - is_openrouter = isinstance(model, str) and model.startswith("openrouter/") - if self.azure and not is_databricks and not is_openrouter: - model = 'azure/' + model - if 'claude' in model and not system: - system = "No system prompt provided" + # Azure mode rewrites only OpenAI models. Explicit non-OpenAI provider + # prefixes must remain intact in multi-provider configurations. + model = completion_model + openrouter_model = self._canonical_openrouter_model(model, request_provider) + normalized_system, user = self.normalize_request_prompts(model, system, user) + if normalized_system != system: get_logger().warning( "Empty system prompt for claude model. Adding a newline character to prevent OpenAI API error.") + system = normalized_system messages = [{"role": "system", "content": system}, {"role": "user", "content": user}] if img_path: - try: - # check if the image link is alive - r = await asyncio.to_thread( - requests.head, - img_path, - allow_redirects=True, - timeout=_IMAGE_HEAD_TIMEOUT_SECONDS, - ) - if r.status_code == 404: - error_msg = "The image link is not [alive](img_path).\nPlease repost the original image as a comment, and send the question again with 'quote reply' (see [instructions](https://pr-agent-docs.codium.ai/tools/ask/#ask-on-images-using-the-pr-code-as-context))." - get_logger().error(error_msg) - return f"{error_msg}", "error" - except Exception as e: - get_logger().error(f"Error fetching image: {img_path}", e) - return f"Error fetching image: {img_path}", "error" messages[1]["content"] = [{"type": "text", "text": messages[1]["content"]}, {"type": "image_url", "image_url": {"url": img_path}}] thinking_kwargs_gpt5 = None + openrouter_reasoning_effort = None # Detect GPT-5 family regardless of provider prefix(es) on the model name. # Users sometimes put a provider prefix in config (e.g. "openai/gpt-5.1-codex-max"), # and Azure mode auto-prepends "azure/", which together can produce stacked prefixes # like "azure/openai/gpt-5...". Without normalization the GPT-5 path is skipped and # litellm rejects the request with UnsupportedParamsError for temperature=0.2. - model_base = model - while model_base.startswith(('openai/', 'azure/')): - model_base = model_base.removeprefix('openai/').removeprefix('azure/') - is_gpt6_astra = model_base.removesuffix('_thinking') == 'gpt-6-astra' - if model_base.startswith('gpt-5') or is_gpt6_astra: + is_gpt6_astra = self._is_gpt6_astra_model(model) + is_gpt5_model = self._is_gpt5_model(openrouter_model or model) + if is_gpt5_model or is_gpt6_astra: # Use configured reasoning_effort or default to MEDIUM - config_effort = get_settings().config.reasoning_effort + config_effort = self._default_reasoning_effort try: ReasoningEffort(config_effort) effort = config_effort @@ -838,46 +3944,49 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: if is_gpt6_astra and effort in (ReasoningEffort.NONE.value, ReasoningEffort.MINIMAL.value): get_logger().info(f"GPT-6 Astra does not support reasoning_effort='{effort}'; using 'low'") effort = ReasoningEffort.LOW.value + elif not is_gpt6_astra and effort == ReasoningEffort.MAX.value: + # 'max' is this project's own alias for "the most reasoning available", + # already translated on the Grok and OpenRouter paths. GPT-5.2 and later + # name that level 'xhigh'; litellm reports supports_xhigh_reasoning_effort + # false for gpt-5 and gpt-5.1, so those stay unsupported either way. + # GPT-6 Astra accepts 'max' natively and is left untouched. + get_logger().info( + "GPT-5 models name their top reasoning level 'xhigh'; " + "using 'xhigh' for reasoning_effort='max'" + ) + effort = ReasoningEffort.XHIGH.value - thinking_kwargs_gpt5 = { - "reasoning_effort": effort, - "allowed_openai_params": ["reasoning_effort"], - } + if openrouter_model: + openrouter_reasoning_effort = effort + else: + thinking_kwargs_gpt5 = { + "reasoning_effort": effort, + "allowed_openai_params": ["reasoning_effort"], + } model_family = "GPT-6 Astra" if is_gpt6_astra else "GPT-5" get_logger().info(f"Using reasoning_effort='{effort}' for {model_family} model") - # Routing priority: Azure mode > explicit provider prefix in user config > openai/ - # default. This preserves an explicit "azure/" the user wrote in config even when - # self.azure is false, and avoids stacking when self.azure already added "azure/". - if self.azure: - provider_prefix = 'azure/' - elif user_model.startswith('azure/'): - provider_prefix = 'azure/' - elif user_model.startswith('openai/'): - provider_prefix = 'openai/' - else: - provider_prefix = 'openai/' - model = provider_prefix + model_base.replace('_thinking', '') # remove _thinking suffix - - # Currently, some models do not support a separate system and user prompts if model in self.user_message_only_models or get_settings().config.custom_reasoning_model: user = f"{system}\n\n\n{user}" system = "" get_logger().info(f"Using model {model}, combining system and user prompts") - messages = [{"role": "user", "content": user}] + if img_path: + content = [{"type": "text", "text": user}, + {"type": "image_url", "image_url": {"url": img_path}}] + else: + content = user + messages = [{"role": "user", "content": content}] - # Build request kwargs after normalizing messages for the target model. - # Databricks selects its endpoint via the DATABRICKS_API_BASE env var; don't let an - # api_base configured by another provider (OpenRouter/Ollama/Azure AD/OpenAI) during - # __init__ override it in multi-provider configs. None lets LiteLLM read the env var. - api_base = os.environ.get("DATABRICKS_API_BASE") if is_databricks else self.api_base + # Build request kwargs after normalizing the model and messages so credentials and + # endpoints can be selected for the provider that will actually receive this call. kwargs = { - "model": model, - "deployment_id": deployment_id, - "messages": messages, - "timeout": get_settings().config.ai_timeout, - "api_base": api_base, - } + "model": model, + "messages": messages, + "timeout": get_settings().config.ai_timeout, + } + if deployment_id: + kwargs["deployment_id"] = deployment_id + kwargs.update(provider_request_params) # Caps the completion client's own per-call retries, which otherwise # multiply this handler's retry attempts. Parsed before the request @@ -893,14 +4002,10 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: if thinking_kwargs_gpt5: kwargs.update(thinking_kwargs_gpt5) - if 'temperature' in kwargs: - del kwargs['temperature'] + if is_gpt5_model or is_gpt6_astra: + kwargs.pop('temperature', None) - custom_llm_provider = str( - getattr(get_settings().litellm, "custom_llm_provider", "") or "" - ).strip().lower() - openrouter_reasoning_effort = None - reasoning_model = model.rsplit(":", 1)[0] if model.startswith("openrouter/") else model + reasoning_model = openrouter_model.rsplit(":", 1)[0] if openrouter_model else model # Add reasoning_effort if model supports it. Match the bare model # id as well as any provider-prefixed form (e.g. # "openrouter/google/gemini-2.5-pro", "gemini/gemini-2.5-pro"), so a @@ -911,30 +4016,13 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: reasoning_model == m or reasoning_model.endswith("/" + m) for m in self.support_reasoning_models ): - config_effort = get_settings().config.reasoning_effort - try: - ReasoningEffort(config_effort) - reasoning_effort = config_effort - except (ValueError, TypeError): - reasoning_effort = ReasoningEffort.MEDIUM.value - if config_effort is not None: - get_logger().warning( - f"Invalid reasoning_effort '{config_effort}' in config. " - f"Using default '{reasoning_effort}'. Valid values: {[e.value for e in ReasoningEffort]}" - ) + config_effort = self._default_reasoning_effort + reasoning_effort = self._resolve_reasoning_effort(openrouter_model or model, config_effort) - clamped_effort = self._clamp_grok_reasoning_effort(model, reasoning_effort) - if clamped_effort != reasoning_effort: - get_logger().info( - f"Grok model {model} does not support reasoning_effort='{reasoning_effort}'; " - f"using '{clamped_effort}' instead." - ) - reasoning_effort = clamped_effort - - if model.startswith("openrouter/"): - # LiteLLM 1.98.0 rejects top-level reasoning_effort for some - # OpenRouter model IDs it does not mark as reasoning-capable; - # defer to OpenRouter's unified reasoning object below. + if openrouter_model: + # LiteLLM rejects top-level reasoning_effort for some OpenRouter + # model IDs it does not mark as reasoning-capable; defer to + # OpenRouter's unified reasoning object below. openrouter_reasoning_effort = reasoning_effort else: get_logger().info(f"Adding reasoning_effort with value {reasoning_effort} to model {model}.") @@ -947,35 +4035,42 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: ) or [] except Exception: supported_params = [] - # LiteLLM 1.98.0 omits reasoning_effort for grok-build-latest + # LiteLLM may omit reasoning_effort for grok-build-latest # and OpenAI-compatible gateway-prefixed Grok IDs. if "reasoning_effort" not in supported_params: kwargs["allowed_openai_params"] = ["reasoning_effort"] # https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking - if self._is_claude_adaptive_thinking_model(model) and get_settings().config.get( - "enable_claude_adaptive_thinking", False): + adaptive_thinking_enabled = self._claude_thinking_controls["enable_claude_adaptive_thinking"] + extended_thinking_enabled = self._claude_thinking_controls["enable_claude_extended_thinking"] + claude_thinking_mode = self._claude_thinking_mode(model) + if claude_thinking_mode == "adaptive": kwargs = self._configure_claude_adaptive_thinking(model, kwargs) - elif ( - model in self.claude_extended_thinking_models - and get_settings().config.get("enable_claude_extended_thinking", False) - ): - if self._is_claude_adaptive_thinking_model(model): - get_logger().warning( - f"Skipping extended thinking for {model}: adaptive-only models reject " - f"budget_tokens. Enable config.enable_claude_adaptive_thinking instead." + elif claude_thinking_mode == "extended": + kwargs = self._configure_claude_extended_thinking(model, kwargs) + elif claude_thinking_mode == "unsupported_extended": + get_logger().warning( + f"Skipping extended thinking for {model}: adaptive-only models reject " + f"budget_tokens. Enable config.enable_claude_adaptive_thinking instead." + ) + elif adaptive_thinking_enabled or extended_thinking_enabled: + message = ( + f"No thinking configuration applied for model {model}: adaptive thinking " + f"requires a recognized claude 5 model name in the id and extended " + f"thinking requires exact membership in claude_extended_thinking_models." + ) + if "arn:aws:bedrock:" in model: + message += ( + " For a Bedrock inference profile, address the model by name and pass " + "the ARN with litellm.model_id." ) - else: - kwargs = self._configure_claude_extended_thinking(model, kwargs) + get_logger().warning(message) # Optional output token limit; 0 = unset. Without max_tokens some # providers apply a low service-side default (Bedrock Converse: 4096, # which reasoning can fully consume, returning empty content). # setdefault keeps the extended-thinking limit authoritative. - try: - max_output_tokens = int(get_settings().config.get("max_output_tokens", 0)) - except (TypeError, ValueError): - max_output_tokens = 0 + max_output_tokens = self._resolve_output_token_limit(model, openrouter_model) if max_output_tokens > 0: output_limit_param = "max_completion_tokens" if is_gpt6_astra else "max_tokens" kwargs.setdefault(output_limit_param, max_output_tokens) @@ -993,16 +4088,6 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: if self.repetition_penalty: kwargs["repetition_penalty"] = self.repetition_penalty - #Added support for extra_headers while using litellm to call underlying model, via a api management gateway, would allow for passing custom headers for security and authorization - if get_settings().get("LITELLM.EXTRA_HEADERS", None): - try: - litellm_extra_headers = json.loads(get_settings().litellm.extra_headers) - if not isinstance(litellm_extra_headers, dict): - raise ValueError("LITELLM.EXTRA_HEADERS must be a JSON object") - except json.JSONDecodeError as e: - raise ValueError(f"LITELLM.EXTRA_HEADERS contains invalid JSON: {str(e)}") - kwargs["extra_headers"] = litellm_extra_headers - # Support for custom OpenAI body fields (e.g., Flex Processing) kwargs = _process_litellm_extra_body(kwargs) @@ -1021,7 +4106,7 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: supported_params = [] if "user" in supported_params: kwargs["user"] = request_user - elif isinstance(model, str) and model.startswith("openrouter/"): + elif openrouter_model: # LiteLLM's OpenRouter transformation does not forward the # standard "user" parameter; extra_body reaches the # OpenAI-compatible request body verbatim. @@ -1052,132 +4137,20 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: # Classic `bedrock/` calls use model_id for Bedrock Runtime inference profiles. # Bedrock Mantle uses Projects, so `bedrock_mantle/` intentionally omits it. - model_id = get_settings().get("litellm.model_id") - if model_id and 'bedrock/' in model: - kwargs["model_id"] = model_id - get_logger().info(f"Using Bedrock custom inference profile: {model_id}") + bedrock_model_id = getattr(self, "_bedrock_model_id", None) + if bedrock_model_id and request_provider == "bedrock": + kwargs["model_id"] = bedrock_model_id + get_logger().info(f"Using Bedrock custom inference profile: {bedrock_model_id}") # OpenRouter provider routing, reasoning control and output cap. # Registered reasoning models inherit config.reasoning_effort when # no OpenRouter-specific effort or token budget is configured. - if isinstance(model, str) and model.startswith("openrouter/"): - openrouter_settings = get_settings().get("openrouter", {}) - extra_body = kwargs.get("extra_body") or {} - - # Normalize operator-controlled config: Dynaconf/env overrides can - # arrive as strings (AUTO_CAST_FOR_DYNACONF is disabled), so coerce - # defensively instead of trusting the declared types. - def _as_list(value): - if isinstance(value, (list, tuple)): - return [str(v).strip() for v in value if str(v).strip()] - if isinstance(value, str): - return [v.strip() for v in value.split(",") if v.strip()] - return [] - - def _as_bool(value, default=True): - if isinstance(value, bool): - return value - if isinstance(value, str): - return value.strip().lower() in ("1", "true", "yes", "on") - return default - - def _as_int(value): - try: - return int(value) - except (TypeError, ValueError): - return 0 - - provider_only = _as_list(openrouter_settings.get("provider_only", [])) - provider_order = _as_list(openrouter_settings.get("provider_order", [])) - if provider_only: - extra_body.setdefault("provider", {})["only"] = provider_only - elif provider_order: - provider = extra_body.setdefault("provider", {}) - provider["order"] = provider_order - provider["allow_fallbacks"] = _as_bool(openrouter_settings.get("allow_fallbacks", True)) - - reasoning = {} - effective_reasoning_effort = str( - openrouter_settings.get("reasoning_effort", "") or "" - ).strip().lower() - reasoning_max_tokens = _as_int(openrouter_settings.get("reasoning_max_tokens", 0)) - if effective_reasoning_effort: - try: - ReasoningEffort(effective_reasoning_effort) - except (TypeError, ValueError): - get_logger().warning( - f"Ignoring invalid openrouter.reasoning_effort '{effective_reasoning_effort}'. " - f"Valid values: {[effort.value for effort in ReasoningEffort]}." - ) - effective_reasoning_effort = "" - if not effective_reasoning_effort: - if reasoning_max_tokens > 0 and openrouter_reasoning_effort: - if openrouter_reasoning_effort == "none": - get_logger().warning( - f"Ignoring config.reasoning_effort='{openrouter_reasoning_effort}' because " - "openrouter.reasoning_max_tokens takes precedence." - ) - else: - get_logger().info( - "Using openrouter.reasoning_max_tokens over" - f" config.reasoning_effort='{openrouter_reasoning_effort}'." - ) - elif reasoning_max_tokens <= 0: - effective_reasoning_effort = openrouter_reasoning_effort or "" - - if effective_reasoning_effort: - clamped_effort = self._clamp_grok_reasoning_effort(model, effective_reasoning_effort) - if clamped_effort != effective_reasoning_effort: - get_logger().info( - f"Grok model {model} does not support reasoning_effort=" - f"'{effective_reasoning_effort}'; using '{clamped_effort}' instead." - ) - effective_reasoning_effort = clamped_effort - - # Preserve explicit disablement; otherwise keep effort and - # max_tokens mutually exclusive by preferring the token budget. - if effective_reasoning_effort == "none": - if reasoning_max_tokens > 0: - get_logger().warning( - "Ignoring openrouter.reasoning_max_tokens because " - "openrouter.reasoning_effort='none' disables reasoning." - ) - reasoning["enabled"] = False - elif reasoning_max_tokens > 0: - if effective_reasoning_effort: - get_logger().warning( - f"Ignoring openrouter.reasoning_effort='{effective_reasoning_effort}' because " - "openrouter.reasoning_max_tokens takes precedence." - ) - reasoning["max_tokens"] = reasoning_max_tokens - elif effective_reasoning_effort: - # OpenRouter uses xhigh for the max alias; extra_body bypasses - # LiteLLM's OpenRouter parameter mapping. - reasoning["effort"] = ( - "xhigh" if effective_reasoning_effort == "max" else effective_reasoning_effort - ) - if reasoning: - get_logger().info(f"Adding OpenRouter reasoning {reasoning} to model {model}.") - extra_body["reasoning"] = reasoning - - if extra_body: - kwargs["extra_body"] = extra_body - - max_tokens = _as_int(openrouter_settings.get("max_tokens", 0)) - if max_tokens > 0: - existing = _as_int(kwargs.get("max_tokens", 0)) - kwargs["max_tokens"] = min(existing, max_tokens) if existing > 0 else max_tokens - effective_max_tokens = _as_int(kwargs.get("max_tokens", 0)) - effective_reasoning_max_tokens = _as_int(reasoning.get("max_tokens", 0)) - if ( - model.startswith("openrouter/anthropic/") - and effective_reasoning_max_tokens > 0 - and 0 < effective_max_tokens <= effective_reasoning_max_tokens - ): - get_logger().warning( - f"OpenRouter Anthropic max_tokens ({effective_max_tokens}) must be greater than " - f"reasoning_max_tokens ({effective_reasoning_max_tokens}) to leave output headroom." - ) + if openrouter_model: + kwargs = self._apply_openrouter_request_controls( + openrouter_model, + kwargs, + openrouter_reasoning_effort, + ) get_logger().debug("Prompts", artifact={"system": system, "user": user}) @@ -1185,14 +4158,6 @@ def _as_int(value): get_logger().info(f"\nSystem prompt:\n{system}") get_logger().info(f"\nUser prompt:\n{user}") - # Inject api_key to the call. This key is populated during init by providers - # like Groq, SambaNova, XAI, Azure AD, and OpenRouter. Skip if None or placeholder. - # Databricks authenticates via the DATABRICKS_API_KEY/DATABRICKS_API_BASE env vars, - # so don't override it with another provider's key in multi-provider configs. - if (litellm.api_key and litellm.api_key != DUMMY_LITELLM_API_KEY - and not is_databricks): - kwargs["api_key"] = litellm.api_key - # Optional fixed provider override, so a raw hosted model id reaches the # provider unchanged instead of being rewritten by LiteLLM's prefix inference. if custom_llm_provider: @@ -1205,12 +4170,21 @@ def _as_int(value): get_logger().error(f"Rate limit error during LLM inference: {e}") raise except openai.APIError as e: - if _bedrock_imds and not self._aws_imds_fell_back and self._aws_static_creds: - self._activate_static_aws_fallback() - # Retry immediately while still holding the lock so that the - # env-var swap is fully visible to this call. Letting @retry - # handle the retry would release the lock between attempts, - # allowing a concurrent coroutine to overwrite os.environ. + if aws_can_fallback: + if not self._aws_imds_fell_back: + self._activate_static_aws_fallback() + fallback_credentials = dict(self._aws_active_creds) + request_region = kwargs.get("aws_region_name") + for key in AWS_REQUEST_CREDENTIAL_KEYS: + kwargs.pop(key, None) + kwargs.update(fallback_credentials) + if request_region and ( + request_provider == "bedrock_mantle" + or (request_provider == "bedrock" and _get_bedrock_model_region( + kwargs["model"], kwargs.get("model_id"), + )) + ): + kwargs["aws_region_name"] = request_region resp, finish_reason, response_obj = await self._get_completion(**kwargs) else: get_logger().warning(f"Error during LLM inference: {e}") @@ -1241,6 +4215,47 @@ def _as_int(value): return resp, finish_reason + async def probe_completion(self, model: str, *, max_tokens: int = 10, timeout: int = 10, _completion=None) -> None: + """Preserve the single-call health probe using request-local credentials.""" + custom_llm_provider = self._custom_llm_provider + configured_deployment_id = self.deployment_id + routed_model = self._route_model_for_request(model, custom_llm_provider, configured_deployment_id) + request_provider = ( + PROVIDER_SETTING_ALIASES.get(custom_llm_provider, custom_llm_provider) + if custom_llm_provider + else self._resolve_request_provider(routed_model) + ) + deployment_id = ( + configured_deployment_id + if request_provider == "azure" and not routed_model.startswith("azure_text/") + else None + ) + async with self._snapshot_aws_request_credentials(self._should_use_aws_imds(request_provider)) as ( + aws_request_credentials, + _, + ): + kwargs = { + "model": self._normalize_gpt5_model_for_request(routed_model, model, custom_llm_provider), + "messages": [{"role": "system", "content": "Say ping"}], + "max_tokens": max_tokens, + "timeout": timeout, + } + if deployment_id: + kwargs["deployment_id"] = deployment_id + kwargs.update(await self._get_provider_request_params_async( + routed_model, + provider=request_provider, + transport_provider=custom_llm_provider or None, + transport_model=deployment_id or kwargs["model"], + aws_request_credentials=aws_request_credentials, + )) + if custom_llm_provider: + kwargs["custom_llm_provider"] = custom_llm_provider + if self._bedrock_model_id and request_provider == "bedrock": + kwargs["model_id"] = self._bedrock_model_id + kwargs["model"] = normalize_litellm_model(kwargs["model"], custom_llm_provider) + await self._acompletion(_completion=_completion, **kwargs) + async def _get_completion(self, **kwargs): """ Wrapper that automatically handles streaming for required models. @@ -1250,32 +4265,25 @@ async def _get_completion(self, **kwargs): # Double the prefix so LiteLLM strips its provider prefix but preserves # OpenRouter's native router ID; leave other explicit providers unchanged. kwargs["model"] = normalize_litellm_model(model, custom_llm_provider) - api_base_value = kwargs.get("api_base") - api_base = api_base_value.strip().lower() if isinstance(api_base_value, str) else "" - force_streaming = ( - bool(custom_llm_provider) - and custom_llm_provider == self.force_streaming_provider - and bool(self.force_streaming_api_base_substrings) - and any(substring in api_base for substring in self.force_streaming_api_base_substrings) - ) + force_streaming = self._force_streaming_for_request(custom_llm_provider, kwargs.get("api_base")) # Some OpenAI-compatible endpoints can return an empty-string # finish_reason on non-streaming responses, which LiteLLM rejects during # response normalization. Streaming avoids that conversion path. - if model in self.streaming_required_models or force_streaming: + if self._requires_streaming(model) or force_streaming: kwargs["stream"] = True kwargs["stream_options"] = {"include_usage": True} - if force_streaming and model not in self.streaming_required_models: + if force_streaming and not self._requires_streaming(model): get_logger().info( f"Using streaming mode for model {model} " "due to OpenAI-compatible endpoint compatibility" ) else: get_logger().info(f"Using streaming mode for model {model}") - response = await acompletion(**kwargs) + response = await self._acompletion(**kwargs) return await _handle_streaming_response(response, model=model) else: - response = await acompletion(**kwargs) + response = await self._acompletion(**kwargs) if response is None or len(response["choices"]) == 0: raise openai.APIError( f"No choices in model response from {model}", @@ -1293,3 +4301,254 @@ async def _get_completion(self, **kwargs): body=None, ) return content, finish_reason, response + + async def _acompletion( + self, _completion=None, _azure_oidc_guard=False, _azure_ad_guard=False, _raw_api_key_guard=False, **kwargs, + ): + """Call LiteLLM with any provider compatibility context scoped to this task.""" + _completion = _completion or acompletion + custom_llm_provider = str(kwargs.get("custom_llm_provider") or "").strip().lower() + provider = ( + PROVIDER_SETTING_ALIASES.get(custom_llm_provider, custom_llm_provider) + if custom_llm_provider + else self._resolve_request_provider(kwargs.get("model")) + ) + transport = ( + custom_llm_provider or self._resolve_request_transport_provider(kwargs.get("model")) or provider + ) + transport_model = kwargs.get("deployment_id") or kwargs.get("model") + azure_ad_token = kwargs.get("azure_ad_token") + oidc_selector = azure_ad_token if isinstance(azure_ad_token, str) and azure_ad_token.startswith("oidc/") else None + companion_auth = ( + self._uses_captured_azure_companion_auth(provider) + and not _is_cloudflare_gateway(kwargs.get("api_base")) + ) + if provider == "azure_ai" and companion_auth: + native_transport, native_model = _azure_ai_native_transport(transport, transport_model) + if native_transport == "azure": + # Keep Azure AI's captured credentials and public kwargs; only + # select the bridge matching LiteLLM's downstream transport. + provider = transport = native_transport + transport_model = native_model + _azure_ad_guard = _azure_ad_guard or _raw_api_key_guard + ordinary_sdk_auth = ( + provider == "azure" and (azure_ad_token or companion_auth) and not oidc_selector + and not _is_cloudflare_gateway(kwargs.get("api_base")) + and _request_local_openai_headers(transport, model=transport_model) is not None + ) + if provider == "azure" and azure_ad_token and kwargs.get("api_key") is not None: + captured_key = self._captured_api_key("azure") + model = kwargs.get("model", "") + headers = dict(kwargs.get("headers") or {}) + guard_key = kwargs.get("api_key") == DUMMY_LITELLM_API_KEY and (self._azure_ad or not captured_key) + cloudflare_key_branch = _is_cloudflare_gateway(kwargs.get("api_base")) + if ( + (guard_key or cloudflare_key_branch) + and (not oidc_selector or not _azure_oidc_guard or cloudflare_key_branch) + and not (ordinary_sdk_auth and _azure_ad_guard) + and (cloudflare_key_branch or (transport != "azure_text" and not model.startswith("azure_text/"))) + and not _uses_openai_responses_transport(transport_model, transport) + ): + for canonical in ("Authorization", "api-key"): + matches = [name for name in headers if name.lower() == canonical.lower()] + if matches: + value = headers[matches[-1]] + for name in matches: + del headers[name] + headers[canonical] = value + if cloudflare_key_branch and (not guard_key or "api-key" in headers): + # Preserve the initial SDK alias/key choice, including its + # absent header; a later alias must not replace this identity. + sdk_token = self._azure_sdk_ad_token + headers.setdefault("Authorization", f"Bearer {sdk_token}" if sdk_token else openai.Omit()) + headers.setdefault("api-key", openai.Omit() if sdk_token else kwargs["api_key"]) + elif guard_key: + # Azure v1 and Cloudflare can prefer the guard key over AD + # auth. Keep the guard, but choose auth through SDK headers. + if "Authorization" not in headers: + if azure_ad_token.startswith("oidc/"): + # Keep LiteLLM's pinned OIDC exchange/cache semantics; + # the selector must never become a Bearer credential. + azure_ad_token = await asyncio.to_thread( + _exchange_azure_oidc_token, + azure_ad_token, + dict(self._azure_oidc_environment), + ) + kwargs["azure_ad_token"] = azure_ad_token + headers["Authorization"] = f"Bearer {azure_ad_token}" + headers.setdefault("api-key", openai.Omit()) + kwargs["headers"] = headers + anthropic_token = None + vertex_token = None + vertex_active_token = None + vertex_aws_token = None + vertex_default_token = None + databricks_token = None + if provider == "vertex_ai": + _install_vertex_executable_guard() + adc_snapshot = self._vertex_gac_adc or self._vertex_default_adc + if adc_snapshot is not None: + _install_vertex_default_adc_bridge() + kwargs.setdefault("vertex_credentials", adc_snapshot["cache_key"]) + vertex_credentials = kwargs.get("vertex_credentials") + if adc_snapshot is not None and vertex_credentials == adc_snapshot["cache_key"]: + vertex_credentials = adc_snapshot["info"] + if isinstance(vertex_credentials, str): + try: + vertex_credentials = json.loads(vertex_credentials) + except (ValueError, TypeError): + vertex_credentials = None + if isinstance(vertex_credentials, dict): + source = vertex_credentials.get("credential_source") + if ( + vertex_credentials.get("type") == "external_account" + and isinstance(source, dict) + and "executable" in source + ): + raise ValueError("Vertex executable credentials are incompatible with request isolation") + _install_vertex_impersonated_credentials_bridge() + if vertex_credentials.get("type") == "external_account": + _install_vertex_wif_project_bridge() + vertex_token = _vertex_request_credentials.set(vertex_credentials) + vertex_active_token = _vertex_request_active.set(True) + vertex_aws_token = _vertex_request_aws_environment.set(self._vertex_aws_environment) + vertex_default_token = _vertex_request_default_adc.set(adc_snapshot) + if provider == "databricks": + captured_key = self._captured_api_key("databricks") + keyless = kwargs.get("api_key") == DUMMY_LITELLM_API_KEY and not captured_key + if keyless: + _install_databricks_keyless_bridge() + databricks_token = _databricks_request_keyless.set(keyless) + if provider == "anthropic": + _install_anthropic_auth_token_bridge() + captured_key = self._captured_api_key("anthropic") + anthropic_token = _anthropic_request_auth_token.set({ + "auth_token": getattr(self, "_anthropic_auth_token", None), + "generated_guard": kwargs.get("api_key") == DUMMY_LITELLM_API_KEY and not captured_key, + }) + sdk_token = None + if _request_local_openai_headers(transport, model=transport_model) is not None: + if kwargs.get("client") is not None: + raise ValueError("Externally supplied SDK clients cannot guarantee request header isolation") + _install_sdk_header_bridge() + explicit_headers = _merge_sdk_headers(kwargs.get("headers") or {}, kwargs.get("extra_headers") or {}) + _check_sdk_marker_collision(explicit_headers) + sdk_token = _sdk_request_headers.set({ + **self._sdk_header_defaults, + "provider": provider, + "explicit_headers": explicit_headers, + }) + oidc_context = None + ad_responses_context = None + raw_key_context = None + raw_auth_context = None + try: + # A tokenless context must not match an unrelated nested None token. + oidc_context = _azure_oidc_request.set(None) + ad_responses_context = _azure_ad_responses_request.set(None) + raw_key_provider = None + if provider == "azure_ai" and companion_auth and _raw_api_key_guard: + _install_raw_api_key_guard_override_bridge(transport) + raw_key_provider = transport + if _raw_api_key_guard and any( + name.lower() == "authorization" for name in (kwargs.get("headers") or {}) + ): + _install_raw_api_key_guard_bridge() + _install_raw_api_key_guard_override_bridge(transport) + raw_key_provider = transport + snapshot = getattr(self, "_raw_guard_auth_snapshot", {}) + if ( + kwargs.get("api_key") == DUMMY_LITELLM_API_KEY + and not snapshot.get("generic_key", True) + and JSONProviderRegistry.supports_responses_api(transport) + and _uses_openai_responses_transport(transport_model, transport) + ): + headers = dict(kwargs["headers"]) + authorization = [name for name in headers if name.lower() == "authorization"] + if len(authorization) == 1: + # Native JSON Responses merges explicit headers after + # creating its guard header, using case-sensitive keys. + headers["Authorization"] = headers.pop(authorization[0]) + kwargs["headers"] = headers + # A nested non-guard call must not inherit an outer guard's provenance. + raw_key_context = _raw_api_key_guard_provider.set(raw_key_provider) + raw_auth_context = _raw_api_key_guard_auth.set( + getattr(self, "_raw_guard_auth_snapshot", None) if raw_key_provider else None + ) + if ( + provider == "azure" and (azure_ad_token or companion_auth) and not oidc_selector + and _uses_openai_responses_transport(transport_model, transport) + ) or (provider == "azure_ai" and companion_auth): + _install_azure_oidc_bridge() + environment = dict(self._azure_oidc_environment) + _azure_ad_responses_request.set({ + "dispatch_token": azure_ad_token, + "generated_guard": _azure_ad_guard or (provider == "azure_ai" and _raw_api_key_guard), + "companion": companion_auth, + "environment": environment, + "auth_environment": {**environment, **self._azure_oidc_auth_environment}, + }) + if provider == "azure" and (oidc_selector or ordinary_sdk_auth): + _install_azure_oidc_bridge() + environment = dict(self._azure_oidc_environment) + if set(environment) != set(AZURE_OIDC_ENV_VARS): + raise RuntimeError("Azure OIDC exchange snapshot is incomplete") + auth_environment = {**environment, **self._azure_oidc_auth_environment} + # A missing selector scopes ordinary AD to SDK identity/cache + # handling without entering the OIDC exchange or resolver. + _azure_oidc_request.set({ + "selector": oidc_selector, + "dispatch_token": azure_ad_token, + "generated_guard": _azure_oidc_guard or _azure_ad_guard, + "environment": environment, + "auth_environment": auth_environment, + "selector_hash": hashlib.sha256(azure_ad_token.encode()).hexdigest() + if isinstance(azure_ad_token, str) else None, + "identity_hash": hashlib.sha256(json.dumps(auth_environment, sort_keys=True).encode()).hexdigest(), + }) + if provider != "bedrock_mantle": + return await _completion(**kwargs) + request_credentials = { + key: kwargs[key] + for key in BEDROCK_MANTLE_REQUEST_CONTEXT_KEYS + if key in kwargs + } + guard_generic_api_key = kwargs.get("api_key") == DUMMY_LITELLM_API_KEY + if not request_credentials and not guard_generic_api_key: + return await _completion(**kwargs) + if BedrockMantleAuthMixin is None: + raise RuntimeError("The installed LiteLLM version does not provide the Bedrock Mantle signer bridge") + # LiteLLM stores aws_* completion kwargs in litellm_params, while the + # Bedrock Mantle signer only receives optional_params. Install the bridge + # only when this provider is used, then scope credentials to this task. + _install_bedrock_mantle_signer_bridge() + credentials_token = _bedrock_mantle_request_credentials.set(request_credentials) + block_bearer_token = _bedrock_mantle_block_bearer.set(guard_generic_api_key) + try: + return await _completion(**kwargs) + finally: + _bedrock_mantle_block_bearer.reset(block_bearer_token) + _bedrock_mantle_request_credentials.reset(credentials_token) + finally: + if raw_auth_context is not None: + _raw_api_key_guard_auth.reset(raw_auth_context) + if raw_key_context is not None: + _raw_api_key_guard_provider.reset(raw_key_context) + if ad_responses_context is not None: + _azure_ad_responses_request.reset(ad_responses_context) + if oidc_context is not None: + _azure_oidc_request.reset(oidc_context) + if sdk_token is not None: + _sdk_request_headers.reset(sdk_token) + if databricks_token is not None: + _databricks_request_keyless.reset(databricks_token) + if vertex_default_token is not None: + _vertex_request_default_adc.reset(vertex_default_token) + if vertex_aws_token is not None: + _vertex_request_aws_environment.reset(vertex_aws_token) + if vertex_active_token is not None: + _vertex_request_active.reset(vertex_active_token) + if vertex_token is not None: + _vertex_request_credentials.reset(vertex_token) + if anthropic_token is not None: + _anthropic_request_auth_token.reset(anthropic_token) diff --git a/pr_agent/algo/ai_handlers/litellm_helpers.py b/pr_agent/algo/ai_handlers/litellm_helpers.py index 48446dbb62..569049d106 100644 --- a/pr_agent/algo/ai_handlers/litellm_helpers.py +++ b/pr_agent/algo/ai_handlers/litellm_helpers.py @@ -152,24 +152,31 @@ def get_repetition_penalty(): return penalty -def _get_azure_ad_token(): +def _get_azure_ad_credential(settings): + """Create an Azure AD credential for one handler/request context.""" + from azure.identity import ClientSecretCredential + + return ClientSecretCredential( + tenant_id=settings.azure_ad.tenant_id, + client_id=settings.azure_ad.client_id, + client_secret=settings.azure_ad.client_secret, + ) + + +def _get_azure_ad_token(credential): """ Generates an access token using Azure AD credentials from settings. Returns: str: The access token """ - from azure.identity import ClientSecretCredential + if credential is None: + raise ValueError("Azure AD credential is required for request-local token resolution") try: - credential = ClientSecretCredential( - tenant_id=get_settings().azure_ad.tenant_id, - client_id=get_settings().azure_ad.client_id, - client_secret=get_settings().azure_ad.client_secret - ) # Get token for Azure OpenAI service token = credential.get_token("https://cognitiveservices.azure.com/.default") return token.token except Exception as e: - get_logger().error(f"Failed to get Azure AD token: {e}") + get_logger().error(f"Failed to get Azure AD token: {type(e).__name__}") raise diff --git a/pr_agent/algo/cli_args.py b/pr_agent/algo/cli_args.py index e5551efbfb..cceac7bdfe 100644 --- a/pr_agent/algo/cli_args.py +++ b/pr_agent/algo/cli_args.py @@ -1,16 +1,22 @@ from base64 import b64decode -from pr_agent.config_security import REPO_HOST_ONLY_KEYS_BY_SECTION, REPO_OVERRIDABLE_KEYS_BY_HOST_SECTION +from pr_agent.config_security import ( + CLI_HOST_ONLY_KEYS_BY_SECTION, + REPO_HOST_ONLY_KEYS_BY_SECTION, + REPO_OVERRIDABLE_KEYS_BY_HOST_SECTION, +) class CliArgs: @staticmethod def _host_only_setting_arg(arg: str) -> str | None: """Return a protected setting token when a CLI arg targets a host-only key.""" - setting_name = arg.removeprefix('--').split('=', 1)[0].replace('__', '.') + setting_name = arg.lstrip('-').split('=', 1)[0].strip().replace('__', '.') section, separator, key = setting_name.partition('.') if not separator: - if section in REPO_OVERRIDABLE_KEYS_BY_HOST_SECTION or section in REPO_HOST_ONLY_KEYS_BY_SECTION: + if (section in REPO_OVERRIDABLE_KEYS_BY_HOST_SECTION + or section in REPO_HOST_ONLY_KEYS_BY_SECTION + or section in CLI_HOST_ONLY_KEYS_BY_SECTION): return f'.{section}' return None @@ -20,6 +26,9 @@ def _host_only_setting_arg(arg: str) -> str | None: host_only_keys = REPO_HOST_ONLY_KEYS_BY_SECTION.get(section, frozenset()) if key.split('.', 1)[0] in host_only_keys: return f'.{section}.{key}' + cli_host_only_keys = CLI_HOST_ONLY_KEYS_BY_SECTION.get(section, frozenset()) + if key.split('.', 1)[0] in cli_host_only_keys: + return f'.{section}.{key}' return None @staticmethod diff --git a/pr_agent/algo/language_handler.py b/pr_agent/algo/language_handler.py index dee3b8a481..9d2e0ec9e3 100644 --- a/pr_agent/algo/language_handler.py +++ b/pr_agent/algo/language_handler.py @@ -77,10 +77,19 @@ def get_language(filename: str) -> str | None: return get_language +def numeric_languages(languages: Dict) -> Dict: + """ + Keep only the {language: size} entries of a provider's get_languages() result. PyGithub 2.x + adds a "url" string to every GET dict it returns, and a str cannot be ranked against sizes. + """ + return {k: v for k, v in (languages or {}).items() if isinstance(v, (int, float))} + + def sort_files_by_main_languages(languages: Dict, files: list): """ Sort files by their main language, put the files that are in the main language first and the rest files after """ + languages = numeric_languages(languages) # sort languages by their size languages_sorted_list = [k for k, v in sorted(languages.items(), key=lambda item: item[1], reverse=True)] # languages_sorted = sorted(languages, key=lambda x: x[1], reverse=True) diff --git a/pr_agent/algo/output_models.py b/pr_agent/algo/output_models.py new file mode 100644 index 0000000000..9f7c47e79a --- /dev/null +++ b/pr_agent/algo/output_models.py @@ -0,0 +1,205 @@ +"""Pydantic models for structured outputs described by prompt templates.""" + +from enum import Enum +from typing import List, Literal, Optional, Union + +from pydantic import BaseModel, ConfigDict, Field, StrictInt, field_validator + + +class SubPR(BaseModel): + model_config = ConfigDict(extra="forbid") + + relevant_files: List[str] + title: str + + +class KeyIssuesComponentLink(BaseModel): + model_config = ConfigDict(extra="forbid") + + relevant_file: str + issue_header: str + issue_content: str + start_line: int + end_line: int + + +class TodoSection(BaseModel): + model_config = ConfigDict(extra="forbid") + + relevant_file: str + line_number: int + content: str + + +class TicketCompliance(BaseModel): + model_config = ConfigDict(extra="forbid") + + ticket_url: str + ticket_requirements: str + fully_compliant_requirements: str + not_compliant_requirements: str + requires_further_human_verification: str + + +class ContributionTimeCostEstimate(BaseModel): + model_config = ConfigDict(extra="forbid") + + best_case: str + average_case: str + worst_case: str + + +class Review(BaseModel): + model_config = ConfigDict(extra="forbid") + + ticket_compliance_check: Optional[List[TicketCompliance]] = None + estimated_effort_to_review: Optional[StrictInt] = Field( + default=None, alias="estimated_effort_to_review_[1-5]", ge=1, le=5 + ) + risk_level: Optional[Literal["low", "medium", "high"]] = None + merge_recommendation: Optional[Literal["safe_to_merge", "merge_with_caution", "changes_required"]] = None + review_priority_files: Optional[List[str]] = None + contribution_time_cost_estimate: Optional[ContributionTimeCostEstimate] = None + score: Optional[StrictInt] = Field(default=None, ge=0, le=100) + relevant_tests: Optional[Literal["Yes", "No"]] = None + insights_from_user_answers: Optional[str] = None + key_issues_to_review: List[KeyIssuesComponentLink] + security_concerns: Optional[str] = None + todo_sections: Optional[Union[List[TodoSection], str]] = None + can_be_split: Optional[List[SubPR]] = Field(default=None, max_length=3) + + @field_validator("risk_level", "merge_recommendation", "relevant_tests", mode="before") + @classmethod + def _strip_literal_values(cls, value): + return value.strip() if isinstance(value, str) else value + + +class PRReview(BaseModel): + model_config = ConfigDict(extra="forbid") + + review: Review + + +class CodeSuggestion(BaseModel): + relevant_file: str + language: str + existing_code: str + suggestion_content: str + improved_code: str + one_sentence_summary: str + label: str + + +class PRCodeSuggestions(BaseModel): + code_suggestions: List[CodeSuggestion] + + +class CodeSuggestionFeedback(BaseModel): + suggestion_summary: str + relevant_file: str + relevant_lines_start: int + relevant_lines_end: int + suggestion_score: int = Field(ge=0, le=10) + why: str + + +class PRCodeSuggestionsFeedback(BaseModel): + code_suggestions: List[CodeSuggestionFeedback] + + +class PRType(str, Enum): + BUG_FIX = "Bug fix" + TESTS = "Tests" + ENHANCEMENT = "Enhancement" + DOCUMENTATION = "Documentation" + OTHER = "Other" + + +class FileDescription(BaseModel): + filename: str + changes_summary: Optional[str] = None + changes_title: str + label: str + + +class PRDescription(BaseModel): + type: List[PRType] = Field(min_length=1) + description: Optional[str] = None + title: str + changes_diagram: Optional[str] = None + pr_files: Optional[List[FileDescription]] = Field(default=None, max_length=20) + + +class PRDescriptionHeaders(BaseModel): + type: List[PRType] = Field(min_length=1) + description: Optional[str] = None + title: str + changes_diagram: Optional[str] = None + + +class PRFilesWalkthrough(BaseModel): + pr_files: List[FileDescription] + + +class Label(str, Enum): + BUG_FIX = "Bug fix" + TESTS = "Tests" + ENHANCEMENT = "Enhancement" + DOCUMENTATION = "Documentation" + OTHER = "Other" + + +class Labels(BaseModel): + labels: List[str] + + +class CodeDocumentationItem(BaseModel): + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + relevant_file: str = Field(alias="relevant file") + relevant_line: StrictInt = Field(alias="relevant line", ge=1) + doc_placement: Literal["before", "after"] = Field(alias="doc placement") + documentation: str + + +class CodeDocumentation(BaseModel): + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + code_documentation: List[CodeDocumentationItem] = Field(alias="Code Documentation") + + @field_validator("code_documentation", mode="after") + @classmethod + def _reject_duplicate_items(cls, items): + serialized = [tuple(sorted(item.model_dump().items())) for item in items] + if len(serialized) != len(set(serialized)): + raise ValueError("Code Documentation items must be unique") + return items + + +class PRRankResponses(BaseModel): + which_response_was_better: Literal[0, 1, 2] + why: str + score_response1: int = Field(ge=1, le=10) + score_response2: int = Field(ge=1, le=10) + + +class RelevantSection(BaseModel): + file_name: str + relevant_section_header_string: str + + +class DocHelper(BaseModel): + user_question: str + response: str + relevant_sections: List[RelevantSection] + question_is_relevant: Optional[int] = Field(default=None, ge=0, le=1) + + +class FileIdxAndPath(BaseModel): + idx: int = Field(ge=0) + file_name: str + + +class DocHeadingsHelper(BaseModel): + user_question: str + relevant_files_ranking: List[FileIdxAndPath] diff --git a/pr_agent/algo/pr_processing.py b/pr_agent/algo/pr_processing.py index cad8513e90..a89c882656 100644 --- a/pr_agent/algo/pr_processing.py +++ b/pr_agent/algo/pr_processing.py @@ -1,6 +1,8 @@ from __future__ import annotations import traceback +from contextvars import ContextVar +from dataclasses import dataclass from typing import Callable, List, Tuple from pr_agent.algo.git_patch_processing import ( @@ -28,6 +30,32 @@ OUTPUT_BUFFER_TOKENS_HARD_THRESHOLD = 1000 MAX_EXTRA_LINES = 10 +_effective_fallback_chain: ContextVar[tuple[tuple[str, str | None], ...] | None] = ContextVar( + "pr_agent_effective_fallback_chain", default=None +) + + +def get_effective_fallback_chain() -> tuple[tuple[str, str | None], ...] | None: + """Return the model/deployment chain selected for the current retry invocation.""" + return _effective_fallback_chain.get() + + +@dataclass +class PreparedPRDiff: + """The single-call diff and compressed file data prepared for one model attempt. + + The compressed file data is request-scoped. It is only reused by a caller that keeps the + same token handler and model, so fallback attempts still rebuild their model-specific budget. + """ + + diff: str + remaining_files_list: list + file_dict: dict | None = None + files_by_name: dict | None = None + model: str | None = None + add_line_numbers_to_hunks: bool = False + token_handler: TokenHandler | None = None + def cap_and_log_extra_lines(value, direction) -> int: try: @@ -47,7 +75,8 @@ def get_pr_diff(git_provider: GitProvider, token_handler: TokenHandler, add_line_numbers_to_hunks: bool = False, disable_extra_lines: bool = False, large_pr_handling=False, - return_remaining_files=False): + return_remaining_files=False, + return_prepared=False): if disable_extra_lines: PATCH_EXTRA_LINES_BEFORE = 0 PATCH_EXTRA_LINES_AFTER = 0 @@ -76,7 +105,12 @@ def get_pr_diff(git_provider: GitProvider, token_handler: TokenHandler, if total_tokens + OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD < get_max_tokens(model): get_logger().info(f"Tokens: {total_tokens}, total tokens under limit: {get_max_tokens(model)}, " f"returning full diff.") - return "\n".join(patches_extended) + full_diff = "\n".join(patches_extended) + if return_prepared: + return PreparedPRDiff(full_diff, [], model=model, + add_line_numbers_to_hunks=add_line_numbers_to_hunks, + token_handler=token_handler) + return full_diff # if we are over the limit, start pruning (If we got here, we will not extend the patches with extra lines) get_logger().info(f"Tokens: {total_tokens}, total tokens over limit: {get_max_tokens(model)}, " @@ -139,6 +173,21 @@ def get_pr_diff(git_provider: GitProvider, token_handler: TokenHandler, get_logger().debug(f"After pruning, added_list_str: {added_list_str}, modified_list_str: {modified_list_str}, " f"deleted_list_str: {deleted_list_str}") + if return_prepared: + files_by_name = { + file.filename: file + for language in pr_languages + for file in language["files"] + } + return PreparedPRDiff( + final_diff, + remaining_files_list, + file_dict, + files_by_name, + model=model, + add_line_numbers_to_hunks=add_line_numbers_to_hunks, + token_handler=token_handler, + ) if not return_remaining_files: return final_diff else: @@ -163,6 +212,107 @@ def get_pr_diff_multiple_patchs(git_provider: GitProvider, token_handler: TokenH return patches_compressed_list, total_tokens_list, deleted_files_list, remaining_files_list, file_dict, files_in_patches_list +def _pack_pr_multi_diffs(file_dict: dict, + token_handler: TokenHandler, + model: str, + max_calls: int, + return_remaining_files: bool): + patches = [] + final_diff_list = [] + files_in_patches = set() + total_tokens = token_handler.prompt_tokens + call_number = 1 + + for filename, data in file_dict.items(): + if call_number > max_calls: + if get_verbosity_level() >= 2: + get_logger().info(f"Reached max calls ({max_calls})") + break + + patch = data["patch"] + new_patch_tokens = data["tokens"] + + if patch and (token_handler.prompt_tokens + new_patch_tokens) > get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD: + if get_settings().config.get("large_patch_policy", "skip") == "skip": + get_logger().warning(f"Patch too large, skipping: {filename}") + continue + if get_settings().config.get("large_patch_policy") == "clip": + delta_tokens = get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD - token_handler.prompt_tokens + patch_clipped = clip_tokens(patch, delta_tokens, delete_last_line=True, + num_input_tokens=new_patch_tokens) + new_patch_tokens = token_handler.count_tokens(patch_clipped) + if patch_clipped and (token_handler.prompt_tokens + new_patch_tokens) > get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD: + get_logger().warning(f"Patch too large, skipping: {filename}") + continue + get_logger().info(f"Clipped large patch for file: {filename}") + patch = patch_clipped + else: + get_logger().warning(f"Patch too large, skipping: {filename}") + continue + + if patch and (total_tokens + new_patch_tokens > get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD): + final_diff_list.append("\n".join(patches)) + patches = [] + total_tokens = token_handler.prompt_tokens + call_number += 1 + if call_number > max_calls: + if get_verbosity_level() >= 2: + get_logger().info(f"Reached max calls ({max_calls})") + break + if get_verbosity_level() >= 2: + get_logger().info(f"Call number: {call_number}") + + if patch: + patches.append(patch) + files_in_patches.add(filename) + total_tokens += new_patch_tokens + if get_verbosity_level() >= 2: + get_logger().info(f"Tokens: {total_tokens}, last filename: {filename}") + + if patches: + final_diff_list.append("\n".join(patches).strip()) + + if not return_remaining_files: + return final_diff_list + + remaining_files_list = [ + filename for filename in file_dict + if filename not in files_in_patches + ] + return final_diff_list, remaining_files_list + + +def _get_pr_multi_diffs_from_prepared(prepared_diff: PreparedPRDiff, + token_handler: TokenHandler, + model: str, + max_calls: int, + return_remaining_files: bool): + """Pack already transformed file patches without repeating preparation work. + + ``get_pr_diff`` and the review chunking path use the same model-specific token handler and + line-number format. Reusing its compressed file dictionary preserves the existing packing + and large-patch policy while avoiding a second provider fetch, patch conversion, and token + count for every file. + """ + file_dict = {} + for filename, data in (prepared_diff.file_dict or {}).items(): + patch = data["patch"] + tokens = data["tokens"] + file = (prepared_diff.files_by_name or {}).get(filename) + if file and file.ai_file_summary and get_settings().get("config.enable_ai_metadata", False): + patch = add_ai_summary_top_patch(file, patch) + tokens = token_handler.count_tokens(patch) + file_dict[filename] = {**data, "patch": patch, "tokens": tokens} + + return _pack_pr_multi_diffs( + file_dict, + token_handler, + model, + max_calls, + return_remaining_files, + ) + + def pr_generate_extended_diff(pr_languages: list, token_handler: TokenHandler, add_line_numbers_to_hunks: bool, @@ -264,7 +414,7 @@ def pr_generate_compressed_diff(top_langs: list, token_handler: TokenHandler, mo # additional iterations (if needed) if large_pr_handling: NUMBER_OF_ALLOWED_ITERATIONS = get_settings().pr_description.get("max_ai_calls", 4) - 1 # one more call is to summarize - for i in range(NUMBER_OF_ALLOWED_ITERATIONS-1): + for _ in range(NUMBER_OF_ALLOWED_ITERATIONS-1): if remaining_files_list: total_tokens, patches, remaining_files_list, files_in_patch_list = generate_full_patch(convert_hunks_to_line_numbers, file_dict, @@ -308,7 +458,7 @@ def generate_full_patch(convert_hunks_to_line_numbers, file_dict, max_tokens_mod patch_final = "" new_patch_tokens = 0 - # If the patch is too large, just show the file name + # If the patch is too large, leave the file in the remaining-files list. if total_tokens + new_patch_tokens > max_tokens_model - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD: # Current logic is to skip the patch if it's too large # TODO: Option for alternative logic to remove hunks from the patch to reduce the number of tokens @@ -335,10 +485,14 @@ async def retry_with_fallback_models(f: Callable, model_type: ModelType = ModelT if routed: # A cheaper primary for a small pull request; config.fallback_models still follow it. all_models[0], all_deployments[0] = routed + # Ignore surplus deployment entries when fewer fallback models are configured; this matches + # the existing retry loop, which stops as soon as the first successful model returns. + effective_chain = tuple(zip(all_models, all_deployments[:len(all_models)], strict=True)) original_deployment_id = get_settings().get("openai.deployment_id", None) + context_token = _effective_fallback_chain.set(effective_chain) try: # try each (model, deployment_id) pair until one is successful, otherwise raise exception - for i, (model, deployment_id) in enumerate(zip(all_models, all_deployments, strict=True)): + for i, (model, deployment_id) in enumerate(effective_chain): try: get_logger().debug( f"Generating prediction with {model}" @@ -357,6 +511,7 @@ async def retry_with_fallback_models(f: Callable, model_type: ModelType = ModelT record_model_used(model, is_fallback=i > 0) return result finally: + _effective_fallback_chain.reset(context_token) get_settings().set("openai.deployment_id", original_deployment_id) @@ -396,7 +551,8 @@ def get_pr_multi_diffs(git_provider: GitProvider, model: str, max_calls: int = 5, add_line_numbers: bool = True, - return_remaining_files: bool = False): + return_remaining_files: bool = False, + prepared_diff: PreparedPRDiff | None = None): """ Retrieves the diff files from a Git provider, sorts them by main language, and generates patches for each file. The patches are split into multiple groups based on the maximum number of tokens allowed for the given model. @@ -405,16 +561,34 @@ def get_pr_multi_diffs(git_provider: GitProvider, git_provider (GitProvider): An object that provides access to Git provider APIs. token_handler (TokenHandler): An object that handles tokens in the context of a pull request. model (str): The name of the model. - max_calls (int, optional): The maximum number of calls to retrieve diff files. Defaults to 5. + max_calls (int, optional): Maximum number of groups for split diffs; the full-diff fast path may still return one group. Defaults to 5. return_remaining_files (bool, optional): Also return the files the token budget left out, in the same shape as `get_pr_diff`. Files without a patch, and delete-only files, are not reported: nothing was omitted for them. Defaults to False. + prepared_diff (PreparedPRDiff, optional): Reuse compressed file data prepared by a preceding + `get_pr_diff` call for the same model attempt. Defaults to None. Returns: List[str]: A list of final diff strings, split into multiple groups based on the maximum number of tokens allowed for the given model. With `return_remaining_files`, a tuple of that list and the list of omitted file names. """ + if ( + prepared_diff is not None + and prepared_diff.file_dict is not None + and prepared_diff.model == model + and add_line_numbers + and prepared_diff.add_line_numbers_to_hunks == add_line_numbers + and prepared_diff.token_handler is token_handler + ): + return _get_pr_multi_diffs_from_prepared( + prepared_diff, + token_handler, + model, + max_calls, + return_remaining_files, + ) + diff_files = git_provider.get_diff_files() # Sort files by main language @@ -426,7 +600,7 @@ def get_pr_multi_diffs(git_provider: GitProvider, PATCH_EXTRA_LINES_BEFORE = cap_and_log_extra_lines(PATCH_EXTRA_LINES_BEFORE, "before") PATCH_EXTRA_LINES_AFTER = cap_and_log_extra_lines(PATCH_EXTRA_LINES_AFTER, "after") - # try first a single run with standard diff string, with patch extension, and no deletions + # First try a single run with the full diff and extended patch context. patches_extended, total_tokens, patches_extended_tokens = pr_generate_extended_diff( pr_languages, token_handler, add_line_numbers_to_hunks=add_line_numbers, @@ -443,17 +617,11 @@ def get_pr_multi_diffs(git_provider: GitProvider, for lang in pr_languages: sorted_files.extend(sorted(lang['files'], key=lambda x: x.tokens, reverse=True)) - patches = [] - final_diff_list = [] - files_in_patches = set() # files that made it into a chunk - total_tokens = token_handler.prompt_tokens - call_number = 1 + # Build the same transformed file dictionary used by the prepared path, preserving the + # descending token order established above. The shared packer then owns chunk boundaries, + # large-patch policy, and remaining-file tracking for both paths. + file_dict = {} for file in sorted_files: - if call_number > max_calls: - if get_verbosity_level() >= 2: - get_logger().info(f"Reached max calls ({max_calls})") - break - original_file_content_str = file.base_file new_file_content_str = file.head_file patch = file.patch @@ -475,66 +643,19 @@ def get_pr_multi_diffs(git_provider: GitProvider, if file.ai_file_summary and get_settings().get("config.enable_ai_metadata", False): patch = add_ai_summary_top_patch(file, patch) new_patch_tokens = token_handler.count_tokens(patch) - - if patch and (token_handler.prompt_tokens + new_patch_tokens) > get_max_tokens( - model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD: - if get_settings().config.get('large_patch_policy', 'skip') == 'skip': - get_logger().warning(f"Patch too large, skipping: {file.filename}") - continue - elif get_settings().config.get('large_patch_policy') == 'clip': - delta_tokens = get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD - token_handler.prompt_tokens - patch_clipped = clip_tokens(patch, delta_tokens, delete_last_line=True, num_input_tokens=new_patch_tokens) - new_patch_tokens = token_handler.count_tokens(patch_clipped) - if patch_clipped and (token_handler.prompt_tokens + new_patch_tokens) > get_max_tokens( - model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD: - get_logger().warning(f"Patch too large, skipping: {file.filename}") - continue - else: - get_logger().info(f"Clipped large patch for file: {file.filename}") - patch = patch_clipped - else: - get_logger().warning(f"Patch too large, skipping: {file.filename}") - continue - - if patch and (total_tokens + new_patch_tokens > get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD): - final_diff = "\n".join(patches) - final_diff_list.append(final_diff) - patches = [] - total_tokens = token_handler.prompt_tokens - call_number += 1 - if call_number > max_calls: # avoid creating new patches - if get_verbosity_level() >= 2: - get_logger().info(f"Reached max calls ({max_calls})") - break - if get_verbosity_level() >= 2: - get_logger().info(f"Call number: {call_number}") - - if patch: - patches.append(patch) - files_in_patches.add(file.filename) - total_tokens += new_patch_tokens - if get_verbosity_level() >= 2: - get_logger().info(f"Tokens: {total_tokens}, last filename: {file.filename}") - - # Add the last chunk - if patches: - final_diff = "\n".join(patches) - final_diff_list.append(final_diff.strip()) - - if not return_remaining_files: - return final_diff_list - - remaining_files_list = [] - for file in sorted_files: - if file.filename in files_in_patches or file.filename in remaining_files_list: - continue - # a file with no patch, or with a delete-only patch, has nothing to review: - # the token budget is not what kept it out of the diff - if not file.patch or handle_patch_deletions(file.patch, file.base_file, file.head_file, - file.filename, file.edit_type) is None: - continue - remaining_files_list.append(file.filename) - return final_diff_list, remaining_files_list + file_dict[file.filename] = { + 'patch': patch, + 'tokens': new_patch_tokens, + 'edit_type': file.edit_type, + } + + return _pack_pr_multi_diffs( + file_dict, + token_handler, + model, + max_calls, + return_remaining_files, + ) def add_ai_metadata_to_diff_files(git_provider, pr_description_files): diff --git a/pr_agent/algo/repo_context.py b/pr_agent/algo/repo_context.py index dbf813b142..12efb8ec76 100644 --- a/pr_agent/algo/repo_context.py +++ b/pr_agent/algo/repo_context.py @@ -14,6 +14,13 @@ REPO_CONTEXT_CACHE_ATTRIBUTE = "_repo_context_cache" REPO_CONTEXT_CACHE_MAX_SIZE = 256 REPO_CONTEXT_CACHE_TTL_SECONDS = 15 * 60 +_DEFAULT_MAX_SIBLING_CONTEXT_FILES = 5 +# Absolute host-controlled ceiling on sibling-repo context fetches per build. The operator-facing +# repo_context_max_sibling_files setting is clamped to this value before any fetch begins, so a +# comment command or an unexpectedly large configured value cannot produce unbounded cross-repo +# API calls. It is intentionally not user-configurable. +_HARD_MAX_SIBLING_CONTEXT_FILES = 20 +_SIBLING_REPO_SEPARATOR = ":" _REPO_CONTEXT_CACHE_MISS = object() _unsupported_repo_context_provider_classes = set() @@ -60,17 +67,17 @@ def _get_markdown_fence(content: str) -> str: def _get_repo_context_cache_key( - context_files: list, max_lines: int, from_default_branch: bool -) -> tuple[tuple[tuple[str, str], ...], int, bool]: + context_files: list, max_lines: int, context_ref: str | None +) -> tuple[tuple[tuple[str, str], ...], int, str | None]: return ( tuple((type(file_path).__name__, str(file_path)) for file_path in context_files), max_lines, - from_default_branch, + context_ref, ) def _get_repo_context_process_cache_key( - git_provider, context_files: list, max_lines: int, from_default_branch: bool + git_provider, context_files: list, max_lines: int, context_ref: str | None ) -> tuple | None: try: pr_url = git_provider.get_pr_url() @@ -81,15 +88,12 @@ def _get_repo_context_process_cache_key( return None return type(git_provider).__name__, pr_url, _get_repo_context_cache_key( - context_files, max_lines, from_default_branch + context_files, max_lines, context_ref ) def _get_repo_context_config() -> tuple[list, int] | None: context_files = get_settings().config.get("repo_context_files", []) - if not context_files: - return None - if isinstance(context_files, str): get_logger().warning( "repo_context_files should be a list of file paths; treating string value as one file path", @@ -103,6 +107,9 @@ def _get_repo_context_config() -> tuple[list, int] | None: ) return None + if not context_files: + return None + max_lines = get_settings().config.get("repo_context_max_lines", 500) try: max_lines = max(0, int(max_lines)) @@ -127,6 +134,51 @@ def _provider_supports_repo_context(git_provider) -> bool: return False +def _provider_supports_sibling_repo_context(git_provider) -> bool: + provider_method = getattr(type(git_provider), "get_sibling_repo_file_content", None) + return provider_method is not None and provider_method is not GitProvider.get_sibling_repo_file_content + + +def _sibling_repo_context_entries(context_files: list) -> list[tuple[str, str]]: + """Return the (repo_id, file_path) pairs that parse as sibling entries. + + Only entries that would really fetch from a sibling repository may force a cache bypass + (build_repo_context), so this parses with the same rules the loader uses: malformed or + non-sibling entries yield no pair and keep the revision-keyed cache usable. Whether the + provider can resolve siblings is decided separately by _provider_supports_sibling_repo_context. + """ + sibling_entries = [] + for entry in context_files: + repo_id, file_path = _parse_repo_context_file_entry(entry) + if repo_id is not None and file_path: + sibling_entries.append((repo_id, file_path)) + return sibling_entries + + +def _parse_repo_context_file_entry(entry: str | dict) -> tuple[str | None, str]: + """Parse a repo-context entry into (sibling repo id, file path). + + A string entry is always a file path in the current repository and yields (None, path). + A sibling entry is a structured ``{"repo_id": ..., "file_path": ...}`` dict naming a + repository in the same namespace/owner (GitLab: ``group/subgroup/project``, GitHub: + ``owner/repo``); its file is read from the sibling's default branch. A malformed entry + yields (None, "") and is skipped by the loader. + """ + if isinstance(entry, dict): + repo_id = entry.get("repo_id") + file_path = entry.get("file_path") + if not isinstance(repo_id, str) or not isinstance(file_path, str): + return None, "" + repo_id = repo_id.strip().strip("/") + file_path = file_path.strip().lstrip("/") + if not repo_id or not file_path: + return None, "" + return repo_id, file_path + if not isinstance(entry, str): + return None, "" + return None, entry.strip() + + def _get_provider_repo_context_cache(git_provider) -> _RepoContextCache: repo_context_cache = getattr(git_provider, REPO_CONTEXT_CACHE_ATTRIBUTE, None) if repo_context_cache is None or not isinstance(repo_context_cache, _RepoContextCache): @@ -136,17 +188,17 @@ def _get_provider_repo_context_cache(git_provider) -> _RepoContextCache: def _get_cached_repo_context( - git_provider, context_files: list, max_lines: int, from_default_branch: bool + git_provider, context_files: list, max_lines: int, context_ref: str | None ): process_cache_key = _get_repo_context_process_cache_key( - git_provider, context_files, max_lines, from_default_branch + git_provider, context_files, max_lines, context_ref ) if process_cache_key is not None: cached_repo_context = _repo_context_process_cache.get(process_cache_key, _REPO_CONTEXT_CACHE_MISS) if cached_repo_context is not _REPO_CONTEXT_CACHE_MISS: return cached_repo_context - cache_key = _get_repo_context_cache_key(context_files, max_lines, from_default_branch) + cache_key = _get_repo_context_cache_key(context_files, max_lines, context_ref) cached_repo_context = _get_provider_repo_context_cache(git_provider).get( cache_key, _REPO_CONTEXT_CACHE_MISS ) @@ -157,13 +209,13 @@ def _get_cached_repo_context( def _store_repo_context( - git_provider, context_files: list, max_lines: int, from_default_branch: bool, repo_context: str + git_provider, context_files: list, max_lines: int, context_ref: str | None, repo_context: str ) -> None: - cache_key = _get_repo_context_cache_key(context_files, max_lines, from_default_branch) + cache_key = _get_repo_context_cache_key(context_files, max_lines, context_ref) _get_provider_repo_context_cache(git_provider)[cache_key] = repo_context process_cache_key = _get_repo_context_process_cache_key( - git_provider, context_files, max_lines, from_default_branch + git_provider, context_files, max_lines, context_ref ) if process_cache_key: _repo_context_process_cache[process_cache_key] = repo_context @@ -185,45 +237,120 @@ def _read_bool_setting(key: str, default: bool) -> bool: return default +def _read_max_sibling_context_files() -> int: + try: + max_siblings = int(get_settings().config.get("repo_context_max_sibling_files", + _DEFAULT_MAX_SIBLING_CONTEXT_FILES)) + except (TypeError, ValueError): + max_siblings = _DEFAULT_MAX_SIBLING_CONTEXT_FILES + # Clamp defensively so a runtime override or an unexpectedly large configured value can never + # exceed the host-controlled hard ceiling, however the value got into the settings. + return min(max(0, max_siblings), _HARD_MAX_SIBLING_CONTEXT_FILES) + + def _load_repo_context_files( git_provider, context_files: list, from_default_branch: bool | None = None -) -> tuple[dict[str, str], bool]: +) -> tuple[list[tuple[str, str]], bool]: if from_default_branch is None: from_default_branch = _read_bool_setting("repo_context_from_default_branch", default=True) - files = {} + # Ordered (label, content) entries rather than a label-keyed dict: a local path can equal a + # sibling's rendered label, and a label-keyed mapping would silently drop one of them. + files = [] had_fetch_error = False - for file_path in context_files: - if not isinstance(file_path, str) or not file_path.strip(): - get_logger().warning("Skipping invalid repo context file path", artifact={"file_path": file_path}) - continue - - file_path = file_path.strip() - try: - content = git_provider.get_repo_file_content(file_path, from_default_branch=from_default_branch) - except Exception as e: - had_fetch_error = True - get_logger().warning(f"Failed to load repo context file: {file_path}", artifact={"error": str(e)}) + max_siblings = _read_max_sibling_context_files() + sibling_fetch_attempts = 0 + seen_sibling_pairs = set() + for entry in context_files: + repo_id, file_path = _parse_repo_context_file_entry(entry) + if not file_path: + if isinstance(entry, str) and not entry.strip(): + get_logger().warning("Skipping invalid repo context file path", artifact={"file_path": entry}) + else: + get_logger().warning( + "Ignoring malformed repo context entry (expected a local file path or a " + "{'repo_id': ..., 'file_path': ...} sibling entry)", + artifact={"entry": entry}, + ) continue + if repo_id is not None: + if not _provider_supports_sibling_repo_context(git_provider): + get_logger().warning( + f"{type(git_provider).__name__} does not support sibling repository context; " + f"skipping {repo_id}:{file_path}" + ) + continue + sibling_pair = (repo_id, file_path) + if sibling_pair in seen_sibling_pairs: + # Duplicate entries must not consume the fetch cap (or the rendered budget): + # fetch each unique sibling file once and let the cap bound the real work. + get_logger().debug(f"Skipping duplicate sibling repo context file: {repo_id}:{file_path}") + continue + seen_sibling_pairs.add(sibling_pair) + if sibling_fetch_attempts >= max_siblings: + get_logger().warning( + f"Stopping repo context at {max_siblings} sibling fetch attempt(s) (" + f"config.repo_context_max_sibling_files); skipping {repo_id}:{file_path}" + ) + continue + sibling_fetch_attempts += 1 + try: + content = git_provider.get_sibling_repo_file_content(repo_id, file_path) + except Exception: + had_fetch_error = True + # A sibling fetch can fail after the provider resolved ids, so keep the full + # traceback; a flat warning hides which request failed and why. + get_logger().exception( + f"Failed to load sibling repo context file: {repo_id}:{file_path}" + ) + continue + # Render the file under its sibling path so the model sees where it came from. + label = f"{repo_id}/{file_path}" + else: + if isinstance(entry, str) and _SIBLING_REPO_SEPARATOR in entry: + # A ':' inside a plain local path used to read as a sibling entry. Structured + # {"repo_id", "file_path"} dicts are the only sibling form now, so hint at the + # migration instead of silently splitting on ':'. + get_logger().warning( + "repo context file path contains ':' and is read as a local file; use a " + "structured {'repo_id': ..., 'file_path': ...} entry in repo_context_files to " + "load a sibling repository file", + artifact={"file_path": entry}, + ) + try: + content = git_provider.get_repo_file_content(file_path, from_default_branch=from_default_branch) + except Exception as e: + had_fetch_error = True + get_logger().warning(f"Failed to load repo context file: {file_path}", artifact={"error": str(e)}) + continue + label = file_path if not content: - get_logger().debug(f"Repo context file is empty or missing: {file_path}") + get_logger().debug(f"Repo context file is empty or missing: {label}") continue if isinstance(content, bytes): content = content.decode("utf-8", errors="replace") - files[file_path] = str(content).rstrip() + files.append((label, str(content).rstrip())) return files, had_fetch_error -def render_instruction_files(files: dict[str, str]) -> str: +def _repo_context_render_entries(files): + """Yield (path, content) pairs, accepting a label-keyed dict or ordered pairs. + + _load_repo_context_files returns ordered (label, content) pairs so a local path and a + sibling's rendered label may coexist; a dict spelling is kept for direct callers.""" + return files.items() if isinstance(files, dict) else files + + +def render_instruction_files(files: dict[str, str] | list[tuple[str, str]]) -> str: parts = [ INSTRUCTION_FILES_INTRO, "", ] - for path, content in files.items(): + for path, content in _repo_context_render_entries(files): scope = path.rsplit("/", 1)[0] if "/" in path else "repo-root" fence = _get_markdown_fence(content) parts.append(f'') @@ -237,7 +364,9 @@ def render_instruction_files(files: dict[str, str]) -> str: return "\n".join(parts) -def render_instruction_files_with_line_budget(files: dict[str, str], max_lines: int) -> str: +def render_instruction_files_with_line_budget( + files: dict[str, str] | list[tuple[str, str]], max_lines: int +) -> str: parts = [ INSTRUCTION_FILES_INTRO, "", @@ -246,7 +375,7 @@ def render_instruction_files_with_line_budget(files: dict[str, str], max_lines: if max_lines < len(parts) + 1: return "" - for path, content in files.items(): + for path, content in _repo_context_render_entries(files): scope = path.rsplit("/", 1)[0] if "/" in path else "repo-root" fence = _get_markdown_fence(content) file_header = [ @@ -289,12 +418,32 @@ def build_repo_context(git_provider) -> str: if not _provider_supports_repo_context(git_provider): return "" - from_default_branch = _read_bool_setting("repo_context_from_default_branch", default=True) - cached_repo_context = _get_cached_repo_context( - git_provider, context_files, max_lines, from_default_branch + # Sibling repositories change independently of the requesting repo, so the current-repo + # revision cannot reconstruct their default-branch content. Bypass the revision-keyed cache + # only when the build can actually attempt a sibling fetch. A zero effective sibling-fetch + # cap disables the fetches entirely, so a local + sibling config with cap 0 must keep the + # revision-keyed cache usable (otherwise disabled siblings would silently disable local + # caching too). Malformed entries and providers without sibling support likewise keep the + # cache usable. + has_sibling_entries = ( + bool(_sibling_repo_context_entries(context_files)) + and _provider_supports_sibling_repo_context(git_provider) + and _read_max_sibling_context_files() > 0 ) - if cached_repo_context is not _REPO_CONTEXT_CACHE_MISS: - return cached_repo_context + + from_default_branch = _read_bool_setting("repo_context_from_default_branch", default=True) + context_ref = None + if not has_sibling_entries: + # Resolve the revision being read once and key the cache on it: within the TTL a rebase + # or a push to the base branch must not serve file content from a commit that has moved. + # A sibling-only build never needs this lookup, whose failure would otherwise abort the + # whole context build before any cross-repository file is loaded. + context_ref = git_provider.get_repo_context_ref(from_default_branch) + cached_repo_context = _get_cached_repo_context( + git_provider, context_files, max_lines, context_ref + ) + if cached_repo_context is not _REPO_CONTEXT_CACHE_MISS: + return cached_repo_context files, had_fetch_error = _load_repo_context_files(git_provider, context_files, from_default_branch) @@ -302,6 +451,6 @@ def build_repo_context(git_provider) -> str: # Only cache when every file was fetched successfully. A transient/unexpected fetch error must # not be cached as a real result, so it is retried instead of being served until the TTL expires. - if not had_fetch_error: - _store_repo_context(git_provider, context_files, max_lines, from_default_branch, repo_context) + if not had_fetch_error and not has_sibling_entries: + _store_repo_context(git_provider, context_files, max_lines, context_ref, repo_context) return repo_context diff --git a/pr_agent/algo/review_finding_state.py b/pr_agent/algo/review_finding_state.py index 154cbed16c..768ecafe88 100644 --- a/pr_agent/algo/review_finding_state.py +++ b/pr_agent/algo/review_finding_state.py @@ -155,6 +155,21 @@ def parse_review_state(comment_body: str) -> ParsedReviewState: return ParsedReviewState(state, present=True, valid=True) +def split_review_state_marker(comment_body: str) -> tuple[str, str]: + """Split a comment into its human text and its hidden state marker. + + Returns ``(body, marker)``; ``marker`` is ``""`` when the comment carries no + complete marker. Providers that cap comment length use this to keep the + marker intact while truncating only the human text. + """ + raw = comment_body or "" + match = _STATE_MARKER_RE.search(raw) + if match is None: + return raw, "" + body = (raw[: match.start()] + raw[match.end():]).rstrip() + return body, match.group(0) + + def serialize_review_state(state: Mapping[str, Any]) -> str: """Serialize state deterministically so repeated updates are diffable.""" if not _is_valid_state(state): diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 135e132121..6faa2af88e 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -66,7 +66,7 @@ class TokenHandler: CLAUDE_MODEL = "claude-3-7-sonnet-20250219" CLAUDE_MAX_CONTENT_SIZE = 9_000_000 # Maximum allowed content size (9MB) for Claude API - def __init__(self, pr=None, vars: dict = {}, system="", user="", model=None): + def __init__(self, pr=None, vars: dict | None = None, system="", user="", model=None): """ Initializes the TokenHandler object. @@ -77,6 +77,8 @@ def __init__(self, pr=None, vars: dict = {}, system="", user="", model=None): - user: The user string. - model: Optional model name whose tokenizer should be used. """ + if vars is None: + vars = {} self.encoder = TokenEncoder.get_token_encoder(model) if pr is not None: diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py index 5b0d4cf175..44b7d9b359 100644 --- a/pr_agent/algo/utils.py +++ b/pr_agent/algo/utils.py @@ -2,7 +2,6 @@ import copy import difflib -import hashlib import html import json import os @@ -10,8 +9,6 @@ import string import sys import textwrap -import time -import traceback from datetime import datetime, timezone from decimal import ROUND_HALF_UP, Decimal from enum import Enum @@ -23,7 +20,6 @@ import requests import yaml from pydantic import BaseModel -from starlette_context import context from pr_agent.algo import MAX_TOKENS from pr_agent.algo.git_patch_processing import ( @@ -34,7 +30,7 @@ from pr_agent.algo.run_details import get_run_details from pr_agent.algo.token_handler import TokenEncoder from pr_agent.algo.types import FilePatchInfo -from pr_agent.config_loader import get_settings, get_verbosity_level, global_settings +from pr_agent.config_loader import get_settings, get_verbosity_level from pr_agent.log import get_logger _ENCODED_USER_TEXT_PREFIX = "__pr_agent_encoded_text__:" @@ -223,14 +219,6 @@ class PRDescriptionHeader(str, Enum): FILE_WALKTHROUGH = "File Walkthrough" -def get_setting(key: str) -> Any: - try: - key = key.upper() - return context.get("settings", global_settings).get(key, global_settings.get(key, None)) - except Exception: - return global_settings.get(key, None) - - def as_review_text(value) -> str: """Flatten a review field the model returned as a list or mapping into readable text. @@ -275,18 +263,6 @@ def emphasize_header(text: str, only_markdown=False, reference_link=None) -> str return text -def unique_strings(input_list: List[str]) -> List[str]: - if not input_list or not isinstance(input_list, list): - return input_list - seen = set() - unique_list = [] - for item in input_list: - if item not in seen: - unique_list.append(item) - seen.add(item) - return unique_list - - def _expand_minute_suffix(text: str) -> str: """Replace minute abbreviations like '30m' with '30 minutes'. @@ -502,7 +478,7 @@ def convert_to_markdown_v2(output_data: dict, markdown_text += f"{emoji} Recommended focus areas for review

\n\n" else: markdown_text += f"### {emoji} Recommended focus areas for review\n\n#### \n" - for i, issue in enumerate(issues): + for issue in issues: try: if not issue or not isinstance(issue, dict): continue @@ -705,7 +681,7 @@ def process_can_be_split(emoji, value): markdown_text += f"{emoji} No multiple PR themes\n\n" else: markdown_text += f"{emoji} {key_nice}

\n\n" - for i, split in enumerate(value): + for split in value: title = split.get('title', '') relevant_files = split.get('relevant_files', []) markdown_text += f"
\nSub-PR theme: {title}\n\n" @@ -1007,7 +983,57 @@ def sanitize_yaml_control_chars(text: str, log: bool = True) -> str: return sanitized -def load_yaml(response_text: str, keys_fix_yaml: List[str] = [], first_key="", last_key="") -> dict: +def _looks_like_more_answer(tail: str) -> bool: + """Whether the text after the fence is more of the answer rather than a sign-off. + + Two signals, because each alone has a blind spot: a tail that parses as a mapping or a + list is structured, but one that continues into prose does not parse at all and is + only recognisable from the shape of its first line. + """ + first_line = next((line for line in tail.split('\n') if line.strip()), '') + if re.match(r'^[A-Za-z_][A-Za-z0-9_]*:(\s|$)', first_line): + return True + try: + return isinstance(yaml.safe_load(tail), (dict, list)) + except Exception: + return False + + +def drop_sign_off_after_wrapper_fence(text: str) -> str: + """Drop a closing remark the model added after the wrapper's closing fence. + + The prompts ask for YAML "and nothing else", but the model sometimes signs off + anyway. That either leaves the document unparseable or, for a single block + scalar, parses the fence and the remark into the value. + + No existing fallback recovers it. The one that extracts a fenced block needs + both fences, but most prompts end with an open fence for the model to continue + from, so the reply carries only the closing one. + """ + lines = text.split('\n') + for i in range(len(lines) - 1, -1, -1): + if lines[i].rstrip() != '```': + continue + tail = '\n'.join(lines[i + 1:]) + if not tail.strip(): + return text + if _looks_like_more_answer(tail): + # Dropping it would publish a partial answer, where the parse failure it + # replaces at least triggers a retry. + return text + candidate = '\n'.join(lines[:i]) + try: + if isinstance(yaml.safe_load(candidate), dict): + return candidate + except Exception: + pass + return text + return text + + +def load_yaml(response_text: str, keys_fix_yaml: List[str] | None = None, first_key="", last_key="") -> dict: + if keys_fix_yaml is None: + keys_fix_yaml = [] response_text_original = copy.deepcopy(response_text) response_text = response_text.strip('\n') # strip the fence label only when it is a complete info string, so a key such as @@ -1016,6 +1042,7 @@ def load_yaml(response_text: str, keys_fix_yaml: List[str] = [], first_key="", l if unfenced == response_text: unfenced = response_text.removeprefix('yaml') response_text = unfenced.rstrip() + response_text = drop_sign_off_after_wrapper_fence(response_text) if response_text.split('\n')[-1] == '```': response_text = response_text.removesuffix('```') response_text = sanitize_yaml_control_chars(response_text) @@ -1046,10 +1073,12 @@ def load_yaml(response_text: str, keys_fix_yaml: List[str] = [], first_key="", l def try_fix_yaml(response_text: str, - keys_fix_yaml: List[str] = [], + keys_fix_yaml: List[str] | None = None, first_key="", last_key="", response_text_original="") -> dict: + if keys_fix_yaml is None: + keys_fix_yaml = [] response_text_lines = response_text.split('\n') keys_yaml = ['relevant line:', 'suggestion content:', 'relevant file:', 'existing code:', @@ -1352,7 +1381,7 @@ def _as_int(value, default: int = 0) -> int: return default -def get_max_tokens(model): +def get_max_tokens(model, ignore_max_model_tokens=False): """ Get the maximum number of tokens allowed for a model. logic: @@ -1364,6 +1393,8 @@ def get_max_tokens(model): For all cases, we further limit the number of tokens to 'config.max_model_tokens' if it is set. This aims to improve the algorithmic quality, as the AI model degrades in performance when the input is too long. + Pass ignore_max_model_tokens=True to keep the unreduced value, for sites that deliberately use the + raw model context size rather than the conservative clamp. """ settings = get_settings() custom_max_tokens = _as_int(settings.config.custom_model_max_tokens) @@ -1437,7 +1468,7 @@ def get_max_tokens(model): ) max_model_tokens = _as_int(settings.config.max_model_tokens) if settings.config.max_model_tokens else 0 - if max_model_tokens > 0: + if max_model_tokens > 0 and not ignore_max_model_tokens: max_tokens_model = min(max_model_tokens, max_tokens_model) return max_tokens_model @@ -1579,12 +1610,24 @@ def find_line_number_of_relevant_line_in_file(diff_files: List[FilePatchInfo], delta = 0 start1, size1, start2, size2 = 0, 0, 0, 0 if absolute_position != -1: # matching absolute to relative + skip_hunk = False for i, line in enumerate(patch_lines): # new hunk if line.startswith('@@'): delta = 0 match = re_hunk_header.match(line) - section_header, size1, size2, start1, start2 = extract_hunk_headers(match) + if match: + skip_hunk = False + section_header, size1, size2, start1, start2 = extract_hunk_headers(match) + else: + # combined/merge hunk headers (e.g. '@@@ ... @@@') cannot be anchored, + # so skip the whole hunk instead of crashing + get_logger().warning("Skipping a line that starts with '@@' but is not a " + "unified hunk header", artifact={"line": line}) + skip_hunk = True + continue + elif skip_hunk: + continue elif not line.startswith('-'): delta += 1 @@ -1609,11 +1652,21 @@ def find_line_number_of_relevant_line_in_file(diff_files: List[FilePatchInfo], def scan_patch_lines(is_match): scan_delta = 0 scan_start2 = 0 + skip_hunk = False for i, line in enumerate(patch_lines): if line.startswith('@@'): scan_delta = 0 header_match = re_hunk_header.match(line) - *_, scan_start2 = extract_hunk_headers(header_match) + if header_match: + skip_hunk = False + *_, scan_start2 = extract_hunk_headers(header_match) + else: + skip_hunk = True + get_logger().warning("Skipping a line that starts with '@@' but is not a " + "unified hunk header", artifact={"line": line}) + continue + elif skip_hunk: + continue elif not line.startswith('-'): scan_delta += 1 @@ -1629,11 +1682,21 @@ def scan_patch_lines(is_match): if position == -1 and relevant_line_in_file[0] == '+': no_plus_line = relevant_line_in_file[1:].lstrip() + skip_hunk = False for i, line in enumerate(patch_lines): if line.startswith('@@'): delta = 0 match = re_hunk_header.match(line) - section_header, size1, size2, start1, start2 = extract_hunk_headers(match) + if match: + skip_hunk = False + section_header, size1, size2, start1, start2 = extract_hunk_headers(match) + else: + get_logger().warning("Skipping a line that starts with '@@' but is not a " + "unified hunk header", artifact={"line": line}) + skip_hunk = True + continue + elif skip_hunk: + continue elif not line.startswith('-'): delta += 1 @@ -1645,64 +1708,6 @@ def scan_patch_lines(is_match): break return position, absolute_position -def get_rate_limit_status(github_token) -> dict: - GITHUB_API_URL = get_settings(use_context=False).get("GITHUB.BASE_URL", "https://api.github.com").rstrip("/") # "https://api.github.com" - # GITHUB_API_URL = "https://api.github.com" - RATE_LIMIT_URL = f"{GITHUB_API_URL}/rate_limit" - HEADERS = { - "Accept": "application/vnd.github.v3+json", - "Authorization": f"token {github_token}" - } - - response = requests.get(RATE_LIMIT_URL, headers=HEADERS) - try: - rate_limit_info = response.json() - if rate_limit_info.get('message') == 'Rate limiting is not enabled.': # for github enterprise - return {'resources': {}} - response.raise_for_status() # Check for HTTP errors - except: # retry - time.sleep(0.1) - response = requests.get(RATE_LIMIT_URL, headers=HEADERS) - return response.json() - return rate_limit_info - - -def validate_rate_limit_github(github_token, installation_id=None, threshold=0.1) -> bool: - try: - rate_limit_status = get_rate_limit_status(github_token) - if installation_id: - get_logger().debug(f"installation_id: {installation_id}, Rate limit status: {rate_limit_status['rate']}") - # validate that the rate limit is not exceeded - # validate that the rate limit is not exceeded - for key, value in rate_limit_status['resources'].items(): - if value['remaining'] < value['limit'] * threshold: - get_logger().error(f"key: {key}, value: {value}") - return False - return True - except Exception as e: - get_logger().error(f"Error in rate limit {e}", - artifact={"traceback": traceback.format_exc()}) - return True - - -def validate_and_await_rate_limit(github_token): - try: - rate_limit_status = get_rate_limit_status(github_token) - # validate that the rate limit is not exceeded - for key, value in rate_limit_status['resources'].items(): - if value['remaining'] < value['limit'] // 80: - get_logger().error(f"key: {key}, value: {value}") - sleep_time_sec = value['reset'] - datetime.now().timestamp() - sleep_time_hour = sleep_time_sec / 3600.0 - get_logger().error(f"Rate limit exceeded. Sleeping for {sleep_time_hour} hours") - if sleep_time_sec > 0: - time.sleep(sleep_time_sec + 1) - rate_limit_status = get_rate_limit_status(github_token) - return rate_limit_status - except: - get_logger().error("Error in rate limit") - return None - def github_action_output(output_data: dict, key_name: str): try: @@ -1763,29 +1768,46 @@ def push_outputs(message_type: str, payload: dict | None = None, markdown: str | record["markdown"] = markdown if "stdout" in channels: - print(json.dumps(record, ensure_ascii=False)) + try: + print(json.dumps(record, ensure_ascii=False)) + except Exception as e: + get_logger().warning(f"push_outputs: stdout failed: {type(e).__name__}") if "file" in channels: - file_path = cfg.get('file_path', 'pr-agent-outputs/reviews.jsonl') - folder = os.path.dirname(file_path) - if folder: - os.makedirs(folder, exist_ok=True) - with open(file_path, 'a', encoding='utf-8') as fh: - fh.write(json.dumps(record, ensure_ascii=False) + "\n") + try: + file_path = cfg.get('file_path', 'pr-agent-outputs/reviews.jsonl') + folder = os.path.dirname(file_path) + if folder: + os.makedirs(folder, exist_ok=True) + with open(file_path, 'a', encoding='utf-8') as fh: + fh.write(json.dumps(record, ensure_ascii=False) + "\n") + except Exception as e: + get_logger().warning(f"push_outputs: file failed: {type(e).__name__}") # Local channels first, network last, so a failed POST can't lose a file write. # allow_redirects=False: never follow a redirect from a configured sink to another host. if "webhook" in channels: - webhook_url = _push_outputs_sink_url(cfg, 'webhook_url') - if webhook_url: - requests.post(webhook_url, json=record, timeout=5, allow_redirects=False) + try: + webhook_url = _push_outputs_sink_url(cfg, 'webhook_url') + if webhook_url: + response = requests.post(webhook_url, json=record, timeout=5, allow_redirects=False) + if not 200 <= response.status_code < 300: + get_logger().warning(f"push_outputs: webhook failed with status {response.status_code}") + except Exception as e: + get_logger().warning(f"push_outputs: webhook failed: {type(e).__name__}") # Slack Incoming Webhooks accept {"text": ...} directly, no relay service needed. if "slack" in channels: - slack_webhook_url = _push_outputs_sink_url(cfg, 'slack_webhook_url') - if slack_webhook_url: - text = markdown if markdown is not None else json.dumps(payload or {}, ensure_ascii=False) - requests.post(slack_webhook_url, json={"text": text}, timeout=5, allow_redirects=False) + try: + slack_webhook_url = _push_outputs_sink_url(cfg, 'slack_webhook_url') + if slack_webhook_url: + text = markdown if markdown is not None else json.dumps(payload or {}, ensure_ascii=False) + response = requests.post(slack_webhook_url, json={"text": text}, timeout=5, + allow_redirects=False) + if not 200 <= response.status_code < 300: + get_logger().warning(f"push_outputs: slack failed with status {response.status_code}") + except Exception as e: + get_logger().warning(f"push_outputs: slack failed: {type(e).__name__}") except Exception as e: # Log only the exception type: requests errors embed the (secret-bearing) URL in their text. get_logger().warning(f"push_outputs failed: {type(e).__name__}") @@ -1899,25 +1921,6 @@ def is_value_no(value): return False -def set_pr_string(repo_name, pr_number): - return f"{repo_name}#{pr_number}" - - -def string_to_uniform_number(s: str) -> float: - """ - Convert a string to a uniform number in the range [0, 1]. - The uniform distribution is achieved by the nature of the SHA-256 hash function, which produces a uniformly distributed hash value over its output space. - """ - # Generate a hash of the string - hash_object = hashlib.sha256(s.encode()) - # Convert the hash to an integer - hash_int = int(hash_object.hexdigest(), 16) - # Normalize the integer to the range [0, 1] - max_hash_int = 2 ** 256 - 1 - uniform_number = float(hash_int) / max_hash_int - return uniform_number - - def process_description(description_full: str) -> Tuple[str, List]: if not description_full: return "", [] diff --git a/pr_agent/cli.py b/pr_agent/cli.py index e13c41891c..b28524f5dc 100644 --- a/pr_agent/cli.py +++ b/pr_agent/cli.py @@ -18,6 +18,61 @@ log_level = os.environ.get("LOG_LEVEL", "INFO") setup_logger(log_level) +_PLAIN_DIFF_MARKDOWN_COMMANDS = frozenset({ + "review", "review_pr", "auto_review", + "describe", "describe_pr", + "improve", "improve_code", + "ask", "ask_question", + "config", "settings", "help", +}) +_PLAIN_DIFF_JSON_COMMANDS = frozenset({"review", "review_pr"}) +_OUTPUT_OPTIONS = ("--output", "--json-output") + + +def _resolve_output_option(parser, arg): + option_text = arg.partition("=")[0] + if option_text in _OUTPUT_OPTIONS: + return option_text + if not option_text.startswith("--") or not parser.allow_abbrev: + return None + + # argparse has no public API for resolving one option spelling. Reuse its + # own option table so misplaced options follow the parser's abbreviation + # rules and future ambiguous prefixes are left untouched. + matches = parser._get_option_tuples(option_text) + if len(matches) == 1 and matches[0][1] in _OUTPUT_OPTIONS: + return matches[0][1] + return None + + +def _validate_output_options(parser, args, diff_mode): + for arg in getattr(args, "rest", []): + option = _resolve_output_option(parser, arg) + if option: + parser.error( + f"{option} must appear before the command " + f"(for example: --stdin {option} result review)" + ) + + output = getattr(args, "output", None) + json_output = getattr(args, "json_output", None) + for option, value in (("--output", output), ("--json-output", json_output)): + if value is not None and not value: + parser.error(f"{option} requires a non-empty path") + + command = args.command.lstrip("/").lower() + if output is not None: + if not diff_mode: + parser.error("--output is only supported in plain-diff mode (--stdin or --diff-file)") + if command not in _PLAIN_DIFF_MARKDOWN_COMMANDS: + parser.error(f"--output is not supported for plain-diff command '{command}'") + + if json_output is not None: + if not diff_mode: + parser.error("--json-output is only supported in plain-diff mode (--stdin or --diff-file)") + if command not in _PLAIN_DIFF_JSON_COMMANDS: + parser.error("--json-output is only supported for plain-diff review commands (review or review_pr)") + def set_parser(): parser = argparse.ArgumentParser(description='AI based pull request analyzer', usage= @@ -71,9 +126,11 @@ def set_parser(): parser.add_argument("--stdin", action="store_true", default=False, help="Read a unified diff from stdin (plain-diff local mode)") parser.add_argument("--output", dest="output", type=str, default=None, - help="Write the result to this file (in addition to stdout)") + help=("Write Plain Diff Markdown output to this file " + "(place before the command)")) parser.add_argument("--json-output", dest="json_output", type=str, default=None, - help="Write the parsed review and token usage to this JSON file") + help=("Write a Plain Diff review and token usage to this JSON file " + "(place before the review command)")) parser.add_argument('command', type=str, help='The', choices=commands, default='review') parser.add_argument('rest', nargs=argparse.REMAINDER, default=[]) return parser @@ -85,7 +142,7 @@ def run_command(pr_url, command): args = set_parser().parse_args(run_command_str.split()) # Run the command. Feedback will appear in GitHub PR comments - run(args=args) + return run(args=args) def run(inargs=None, args=None): @@ -93,11 +150,10 @@ def run(inargs=None, args=None): if not args: args = parser.parse_args(inargs) diff_mode = getattr(args, "stdin", False) or getattr(args, "diff_file", None) - if getattr(args, "json_output", None) and not diff_mode: - parser.error("--json-output is only supported in plain-diff mode (--stdin or --diff-file)") + if diff_mode and args.stdin and args.diff_file: + parser.error("--stdin and --diff-file are mutually exclusive") + _validate_output_options(parser, args, diff_mode) if diff_mode: - if args.stdin and args.diff_file: - parser.error("--stdin and --diff-file are mutually exclusive") if args.diff_file: try: with open(args.diff_file, "r", encoding="utf-8") as fh: @@ -122,6 +178,8 @@ def run(inargs=None, args=None): return command = args.command.lower() + settings = get_settings() + propagate_tool_errors_before = settings.config.get("propagate_tool_errors", False) get_settings().set("CONFIG.CLI_MODE", True) # Strip each candidate independently so a whitespace-only CLI value doesn't # short-circuit the PR_AGENT_CONFIG_BRANCH env fallback before precedence. @@ -157,10 +215,15 @@ async def inner(): return result - result = asyncio.run(inner()) - if not result: - parser.print_help() + try: + result = asyncio.run(inner()) + if not result: + parser.print_help() + if result is False and settings.config.get("propagate_tool_errors", False): + return 1 + finally: + settings.set("CONFIG.PROPAGATE_TOOL_ERRORS", propagate_tool_errors_before) if __name__ == '__main__': - run() + raise SystemExit(run()) diff --git a/pr_agent/cli_pip.py b/pr_agent/cli_pip.py index caa56f0c9b..db6670ab08 100644 --- a/pr_agent/cli_pip.py +++ b/pr_agent/cli_pip.py @@ -16,8 +16,8 @@ def main(): get_settings().set("github.user_token", user_token) # Run the command. Feedback will appear in GitHub PR comments - cli.run_command(pr_url, command) + return cli.run_command(pr_url, command) if __name__ == '__main__': - main() + raise SystemExit(main()) diff --git a/pr_agent/config_security.py b/pr_agent/config_security.py index 226092cc58..2dfe762b05 100644 --- a/pr_agent/config_security.py +++ b/pr_agent/config_security.py @@ -30,4 +30,126 @@ # disclosed in a PR comment, so the PR author must not be able to enable it. REPO_HOST_ONLY_KEYS_BY_SECTION = { "pr_reviewer": frozenset({"publish_error_details"}), + # repo_context_sibling_repos lists the sibling repositories whose files a consuming repo + # (or a comment command) may select into model context. A repo's .pr_agent.toml alone must + # not be able to name an arbitrary same-owner private sibling: the actor check bounds who + # triggers the read, not who chose the target or where the output lands, so a sibling + # collaborator running /review on a public repo would otherwise print the sibling's private + # content into that public thread. The list stays host-only (default empty). + # repo_context_max_sibling_files bounds sibling-repository fetches per repo-context build. + # Letting a repo's .pr_agent.toml or a comment command raise it would defeat the safety + # bound and let a commenter force unbounded cross-repository API calls; it stays host-only. + "config": frozenset({"repo_context_max_sibling_files", "repo_context_sibling_repos"}), +} + +# Keys that repositories may still configure from their own default-branch settings but that +# comment/CLI *arguments* must never override. repo_context_files selects which repository and +# sibling files are fetched and rendered as the model's instruction context, so an untrusted +# commenter must not be able to point the bot at arbitrary sibling repo content. Values curated +# by the repo's maintainers in .pr_agent.toml stay accepted (apply_repo_settings does not consult +# this map); only CliArgs.validate_user_args enforces it, so repo settings and comment args do +# not drift. +CLI_HOST_ONLY_KEYS_BY_SECTION = { + "config": frozenset({"repo_context_files"}), +} + +# Keys a per-directory `.pr_agent.toml` can never override, even when their section is +# otherwise open (None) in REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS. Nested files live in +# the working repository where any contributor can edit them, so keys that perform +# host-side writes (label mutation, resolving human review threads, publishing inline +# review findings, replacing pull-request labels with reviewer effort/security labels) or +# consume unbounded external resources (forcing a full issue-index refresh, scanning +# arbitrary issue counts, or repointing the vector backend) stay root-config- or host- +# controlled. Likewise, keys that rewrite pull-request metadata (AI title generation, +# comment-only description publication) stay root-controlled so a nested file cannot +# bypass the operator's choice of which PR fields the bot edits. Ticket extraction +# (require_ticket_analysis_review) stays root-controlled too: enabling it from a nested +# file would make the bot run authenticated Jira lookups the root config already disabled. +# Self-review workflow controls (demand_code_suggestions_self_review and +# approve_pr_on_self_review) stay root-controlled as well: a nested file must not be able +# to make /improve demand a checklist and then auto-approve the pull request when the +# author ticks it. Thread-history collection (pr_questions.use_conversation_history) is +# also root-controlled so a nested file cannot re-enable sending private review-thread +# discussion bodies to the model after the operator opted out. +# Similarly, budget/call-count controls (max_number_of_calls, max_ai_calls, parallel_calls, +# enable_large_pr_chunking, enable_large_pr_handling, async_ai_calls) are restricted so +# a nested file cannot multiply AI calls independently of the host-trusted defaults. +PER_DIRECTORY_HOST_ONLY_KEYS_BY_SECTION = { + "pr_reviewer": frozenset({ + "enable_large_pr_chunking", "max_number_of_calls", + "inline_key_issues", "enable_review_labels_security", + "enable_review_labels_effort", "require_estimate_effort_to_review", + "require_security_review", "require_ticket_analysis_review", + }), + "pr_description": frozenset({ + "publish_labels", "enable_large_pr_handling", "max_ai_calls", "async_ai_calls", + "generate_ai_title", "publish_description_as_comment", + "publish_description_as_comment_persistent", + }), + "pr_questions": frozenset({"resolve_threads", "use_conversation_history"}), + "pr_code_suggestions": frozenset({ + "commitable_code_suggestions", + "max_number_of_calls", "parallel_calls", + "approve_pr_on_self_review", "demand_code_suggestions_self_review", + }), + "pr_similar_issue": frozenset({"force_update_dataset", "max_issues_to_scan", "vectordb", "skip_comments"}), +} + +# Sections a *per-directory* `.pr_agent.toml` may override at all. Nested config +# files live in the working repository where any contributor can edit them, so this +# layer is deliberately narrower than the root/global repo settings: tool +# instructions, suggestion limits, ignore lists and model routing only. A frozenset +# value restricts the section to those keys; None allows every key in the section. +# Secrets, identity and deployment-critical settings (provider tokens, git_provider, +# push_outputs, skills, prompt_fragments, ...) can only be influenced from the root +# config or host environment, never from a nested file. +# +# Tool sections that can trigger bot-side writes (commits, changelog pushes) or +# read/connect from arbitrary URLs or paths are likewise restricted to drop-only keys: +# `pr_update_changelog` cannot be given `push_changelog_changes` (a nested config +# must not cause the bot to commit), and `pr_help_docs` cannot be given `repo_url`, +# `docs_path`, or `supported_doc_exts` (`repo_url` can be resolved into a token-embedded +# clone URL, while `docs_path` plus `supported_doc_exts` can point collection at any +# repository path or extension, so a nested file could read arbitrary source files into +# the model prompt). +# +# The `ignore` section is open to `glob` only: fnmatch translates glob patterns +# into bounded regexes, whereas `ignore.regex` accepts arbitrary expressions that +# filter_ignored() compiles and matches against every changed filename on every +# review. A catastrophic-backtracking pattern committed in a nested file could +# stall a worker, so nested files keep the bounded glob form only. +# +# The `config` section lists model-routing and output knobs but deliberately +# excludes the repo-context builders: `repo_context_files` fetches every listed +# file in full and `repo_context_max_lines` sizes the trimmed output, so without +# a host-trusted upper bound a nested file could balloon repository calls and +# token budget whenever a tool builds context for a directory it crosses. Those +# knobs stay root-/host-controlled. `fallback_models` is excluded too: the retry +# helper treats each entry as one routing attempt per failing model, so an +# arbitrarily long nested list could multiply AI calls; fallback routing stays +# root-/host-controlled. The token-budget keys (`max_model_tokens`, +# `custom_model_max_tokens`, `max_output_tokens`) are excluded as well because +# they directly size request context and completion limits, so without a trusted +# ceiling a nested file could inflate the size and cost of every model request. +REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS = { + "config": frozenset({ + "model", "model_weak", "model_reasoning", + "temperature", "response_language", + "repo_context_from_default_branch", + }), + "ignore": frozenset({"glob"}), + "pr_reviewer": None, + "pr_description": None, + "pr_questions": None, + "pr_code_suggestions": None, + "pr_custom_prompt": None, + "pr_add_docs": None, + "pr_update_changelog": frozenset({"extra_instructions", "add_pr_link"}), + "pr_analyze": None, + "pr_test": None, + "pr_improve_component": None, + "pr_help": None, + "pr_help_docs": frozenset({"exclude_root_readme", "enable_help_text"}), + "pr_similar_issue": None, + "pr_find_similar_component": None, } diff --git a/pr_agent/custom_merge_loader.py b/pr_agent/custom_merge_loader.py index 8dc7d54b89..ad88286a8e 100644 --- a/pr_agent/custom_merge_loader.py +++ b/pr_agent/custom_merge_loader.py @@ -19,10 +19,10 @@ def load(obj, env=None, silent=True, key=None, filename=None): - Supports Dynaconf's fresh_vars feature for dynamic reloading. Args: obj: The Dynaconf settings instance to update. - env: The current environment name (upper case). Defaults to 'DEVELOPMENT'. Note: currently unused. + env: Unused compatibility parameter. Defaults to None. silent (bool): If True, suppress exceptions and log warnings/errors instead. key (str | None): Load only this top-level key (section) if provided; otherwise, load all keys from the files. - filename (str | None): Custom filename for tests (not used when settings_files are provided). + filename (str | None): Unused compatibility parameter. Settings files come from the Dynaconf object. Returns: None """ diff --git a/pr_agent/git_providers/azuredevops_provider.py b/pr_agent/git_providers/azuredevops_provider.py index 67ae2d52e2..6ab1f45a8c 100644 --- a/pr_agent/git_providers/azuredevops_provider.py +++ b/pr_agent/git_providers/azuredevops_provider.py @@ -2,8 +2,8 @@ import datetime as _dt import json -import os import re +from collections import Counter from types import SimpleNamespace from typing import Optional, Tuple from urllib.parse import quote, unquote, urlparse @@ -19,7 +19,7 @@ get_inline_comment_store, has_marker, ) -from ..algo.language_handler import is_valid_file +from ..algo.language_handler import build_language_file_matcher, is_valid_file from ..algo.utils import ( PRCodeSuggestionsIdentity, PRDescriptionHeader, @@ -480,6 +480,9 @@ def supports_line_question_history(self) -> bool: def supports_thread_resolution(self) -> bool: return True + def supports_linked_work_item_tickets(self) -> bool: + return True + def set_pr(self, pr_url: str): self.diff_files = None self._diff_path_map = None @@ -728,6 +731,13 @@ def get_repo_file_content(self, file_path: str, from_default_branch: bool = Fals return "" raise + def get_repo_context_ref(self, from_default_branch: bool = False) -> Optional[str]: + # The PR target (base) commit is the cached revision; the default branch is selected by + # omitting the version, so it carries no explicit ref, mirroring get_repo_file_content. + if from_default_branch: + return None + return self.pr.last_merge_target_commit.commit_id + def get_files(self): if (isinstance(getattr(self, "incremental", None), IncrementalPR) and self.incremental.is_incremental @@ -976,23 +986,6 @@ def publish_comment(self, pr_comment: str, is_temporary: bool = False, thread_co self.temp_comments.append(created_comment) return created_comment - def publish_persistent_comment(self, pr_comment: str, - initial_header: str, - update_header: bool = True, - name='review', - final_update_message=True, - identity_marker: str | None = None, - legacy_initial_header: str | None = None): - return self.publish_persistent_comment_full( - pr_comment, - initial_header, - update_header, - name, - final_update_message, - identity_marker=identity_marker, - legacy_initial_header=legacy_initial_header, - ) - def supports_review_comment_identity(self) -> bool: return True @@ -1167,7 +1160,24 @@ def get_title(self): return self.pr.title def get_languages(self): - languages = [] + # Return {language name: percentage}, like the other providers. Keys are + # language NAMES (e.g. "Python"), not raw extensions: the consumer + # sort_files_by_main_languages() maps each name back to its extensions, so + # returning extensions ("py") silently drops every file into the "Other" + # bucket and defeats the prioritisation. Use the shared configured + # filename matcher so all providers apply the same rules. Percentages are + # computed over the repository's blob inventory (not the PR change set), + # so transient files or untouched files in the PR cannot skew the ranking. + lang_map = get_settings().get("language_extension_map_org", {}) or {} + get_language = build_language_file_matcher(lang_map) + + # Azure may omit merge metadata for a PR; without a target commit there is no + # reliable revision to enumerate, so mirror the empty-map fallback used by the + # other providers for unrecognized repositories. + target_commit = getattr(self.pr, "last_merge_target_commit", None) + if target_commit is None or not getattr(target_commit, "commit_id", None): + return {} + files = self.azure_devops_client.get_items( project=self.workspace_slug, repository_id=self.repo_slug, @@ -1175,25 +1185,21 @@ def get_languages(self): include_content_metadata=True, include_links=False, download=False, + version_descriptor=GitVersionDescriptor( + version=target_commit.commit_id, version_type="commit" + ), ) - for f in files: - if f.git_object_type == "blob": - file_name, file_extension = os.path.splitext(f.path) - languages.append(file_extension[1:]) - - extension_counts = {} - for ext in languages: - if ext != "": - extension_counts[ext] = extension_counts.get(ext, 0) + 1 - total_extensions = sum(extension_counts.values()) - - extension_percentages = { - ext: (count / total_extensions) * 100 - for ext, count in extension_counts.items() - } + lang_count = Counter() + for f in files: + if f.git_object_type != "blob" or not f.path: + continue + language = get_language(f.path) + if language: + lang_count[language] += 1 - return extension_percentages + total = sum(lang_count.values()) or 1 + return {lang: count / total * 100 for lang, count in lang_count.items()} def get_pr_branch(self): pr_info = self.azure_devops_client.get_pull_request_by_id( @@ -1652,9 +1658,6 @@ def get_pr_id(self): get_logger().info(f"Failed to get PR id, error: {e}") return "" - def publish_file_comments(self, file_comments: list) -> bool: - pass - def get_line_link(self, relevant_file: str, relevant_line_start: int, relevant_line_end: int = None) -> str: return self.pr_url+f"?_a=files&path={relevant_file}" diff --git a/pr_agent/git_providers/bitbucket_provider.py b/pr_agent/git_providers/bitbucket_provider.py index a6ae1ec8e6..413548db53 100644 --- a/pr_agent/git_providers/bitbucket_provider.py +++ b/pr_agent/git_providers/bitbucket_provider.py @@ -121,9 +121,12 @@ def _fetch_global_repo_settings(self, workspace): def get_repo_file_content(self, file_path: str, from_default_branch: bool = False): # Read from the PR destination (target) branch, matching the other providers, # or from the repository default branch when from_default_branch is requested. - branch = self.get_repo_default_branch() if from_default_branch else self.pr.destination_branch + branch = self.get_repo_context_ref(from_default_branch) return self.get_pr_file_content(file_path, branch, propagate_errors=True) + def get_repo_context_ref(self, from_default_branch: bool = False) -> Optional[str]: + return self.get_repo_default_branch() if from_default_branch else self.pr.destination_branch + def get_git_repo_url(self, pr_url: str=None) -> str: #bitbucket does not support issue url, so ignore param try: parsed_url = urlparse(self.pr_url) @@ -219,12 +222,8 @@ def publish_code_suggestions(self, code_suggestions: list) -> bool: get_logger().error(f"Bitbucket failed to publish code suggestion, error: {e}") return False - def publish_file_comments(self, file_comments: list) -> bool: - pass - def is_supported(self, capability: str) -> bool: - if capability in ['publish_inline_comments', 'get_labels', - 'gfm_markdown', 'publish_file_comments']: + if capability in ['publish_inline_comments', 'get_labels', 'gfm_markdown']: return False if capability == "push_code" and get_settings().config.restricted_mode: return False @@ -375,6 +374,10 @@ def _get_cloud_comment(self, comment): return comment._cloud_comment return comment + def get_comment_url(self, comment): + comment = self._get_cloud_comment(comment) + return comment.data["links"]["html"]["href"] + def supports_review_comment_identity(self) -> bool: return True @@ -388,28 +391,6 @@ def edit_comment(self, comment, body: str): get_logger().exception(f"Failed to update comment, error: {e}") return False - def publish_persistent_comment( - self, - pr_comment: str, - initial_header: str, - update_header: bool = True, - name='review', - final_update_message=True, - as_thread: bool = False, - identity_marker: str | None = None, - legacy_initial_header: str | None = None, - ): - return self.publish_persistent_comment_full( - pr_comment, - initial_header, - update_header, - name, - final_update_message, - as_thread=as_thread, - identity_marker=identity_marker, - legacy_initial_header=legacy_initial_header, - ) - def publish_comment(self, pr_comment: str, is_temporary: bool = False): if is_temporary and not get_settings().config.publish_output_progress: get_logger().debug(f"Skipping publish_comment for temporary comment: {pr_comment}") @@ -536,8 +517,8 @@ def get_title(self): return self.pr.title def get_languages(self): - languages = {self._get_repo().get_data("language"): 0} - return languages + language = self._get_repo().get_data("language") + return {language: 0} if language else {} def get_pr_branch(self): return self.pr.source_branch diff --git a/pr_agent/git_providers/bitbucket_server_provider.py b/pr_agent/git_providers/bitbucket_server_provider.py index d2ac3b16d7..2c379a7fc0 100644 --- a/pr_agent/git_providers/bitbucket_server_provider.py +++ b/pr_agent/git_providers/bitbucket_server_provider.py @@ -2,6 +2,7 @@ import re import shlex import subprocess +from collections import Counter from types import SimpleNamespace from typing import Optional, Tuple from urllib.parse import quote_plus, urlparse @@ -12,7 +13,7 @@ from ..algo.file_filter import filter_ignored from ..algo.git_patch_processing import decode_if_bytes -from ..algo.language_handler import is_valid_file +from ..algo.language_handler import build_language_file_matcher, is_valid_file from ..algo.types import EDIT_TYPE, FilePatchInfo from ..algo.utils import find_line_number_of_relevant_line_in_file, load_large_diff from ..config_loader import get_settings, get_verbosity_level @@ -138,12 +139,14 @@ def _fetch_global_repo_settings(self, workspace): def get_repo_file_content(self, file_path: str, from_default_branch: bool = False): # Read from the PR target ref (the branch being merged into), matching the other providers, # or from the repository default branch when from_default_branch is requested. + ref = self.get_repo_context_ref(from_default_branch) + return self.get_file(file_path, ref) + + def get_repo_context_ref(self, from_default_branch: bool = False) -> Optional[str]: if from_default_branch: default_branch_dict = self.bitbucket_client.get_default_branch(self.workspace_slug, self.repo_slug) - ref = default_branch_dict.get('displayId') or self.pr.toRef['latestCommit'] - else: - ref = self.pr.toRef['latestCommit'] - return self.get_file(file_path, ref) + return default_branch_dict.get('displayId') or self.pr.toRef['latestCommit'] + return self.pr.toRef['latestCommit'] def get_pr_id(self): return self.pr_num @@ -214,14 +217,15 @@ def publish_code_suggestions(self, code_suggestions: list) -> bool: get_logger().error(f"Failed to publish code suggestion, error: {e}") return False - def publish_file_comments(self, file_comments: list) -> bool: - pass - def is_supported(self, capability: str) -> bool: - if capability in ['get_labels', 'gfm_markdown', 'publish_file_comments']: + if capability in ['get_labels', 'gfm_markdown']: return False return True + def supports_markdown_tables(self) -> bool: + # Bitbucket Data Center renders Markdown tables in comments, but not GFM. + return True + def set_pr(self, pr_url: str): self.workspace_slug, self.repo_slug, self.pr_num = self._parse_pr_url(pr_url) self.pr = self._get_pr() @@ -353,25 +357,6 @@ def publish_comment(self, pr_comment: str, is_temporary: bool = False, as_thread ) return None - def publish_persistent_comment(self, pr_comment: str, - initial_header: str, - update_header: bool = True, - name='review', - final_update_message=True, - as_thread: bool = False, - identity_marker: str | None = None, - legacy_initial_header: str | None = None): - return self.publish_persistent_comment_full( - pr_comment, - initial_header, - update_header, - name, - final_update_message, - as_thread, - identity_marker=identity_marker, - legacy_initial_header=legacy_initial_header, - ) - def supports_review_comment_identity(self) -> bool: return True @@ -426,7 +411,26 @@ def remove_initial_comment(self): get_logger().exception(f"Failed to remove temp comments, error: {e}") def remove_comment(self, comment): - pass + comment_id = comment.get("id") if isinstance(comment, dict) else getattr(comment, "id", None) + comment_version = comment.get("version") if isinstance(comment, dict) else getattr(comment, "version", None) + if not isinstance(comment_id, int) or not isinstance(comment_version, int): + get_logger().warning( + f"Failed to remove Bitbucket Server comment: invalid id or version for {comment!r}" + ) + return False + + try: + self.bitbucket_client.delete_pull_request_comment( + self.workspace_slug, + self.repo_slug, + self.pr_num, + comment_id, + comment_version, + ) + return True + except Exception as e: + get_logger().exception(f"Failed to remove comment, error: {e}") + return False # function to create_inline_comment def create_inline_comment(self, body: str, relevant_file: str, relevant_line_in_file: str, @@ -541,7 +545,20 @@ def get_title(self): return self.pr.title def get_languages(self): - return {"yaml": 0} # devops LOL + # Return {language name: percentage}, like the other providers. + lang_map = get_settings().get("language_extension_map_org", {}) or {} + get_language = build_language_file_matcher(lang_map) + + lang_count = Counter() + for filename in self.get_files(): + if not filename: + continue + language = get_language(filename) + if language: + lang_count[language] += 1 + + total = sum(lang_count.values()) or 1 + return {lang: count / total * 100 for lang, count in lang_count.items()} def get_pr_branch(self): return self.pr.fromRef['displayId'] diff --git a/pr_agent/git_providers/codecommit_client.py b/pr_agent/git_providers/codecommit_client.py index bca9a54bc9..f134acfcb8 100644 --- a/pr_agent/git_providers/codecommit_client.py +++ b/pr_agent/git_providers/codecommit_client.py @@ -54,6 +54,7 @@ class CodeCommitClient: def __init__(self): self.boto_client = None + self.comments_page_size = 100 def is_supported(self, capability: str) -> bool: if capability in ["gfm_markdown"]: @@ -235,7 +236,7 @@ def publish_comment(self, repo_name: str, pr_number: int, destination_commit: st It does not support the ending line number to highlight a range of lines. Returns: - - None + - The boto3 post_comment_for_pull_request response Boto3 Documentation: - aws codecommit post_comment_for_pull_request @@ -248,7 +249,7 @@ def publish_comment(self, repo_name: str, pr_number: int, destination_commit: st # If the comment has code annotations, # then set the file path and line number in the location dictionary if annotation_file and annotation_line: - self.boto_client.post_comment_for_pull_request( + return self.boto_client.post_comment_for_pull_request( pullRequestId=str(pr_number), repositoryName=repo_name, beforeCommitId=destination_commit, @@ -262,7 +263,7 @@ def publish_comment(self, repo_name: str, pr_number: int, destination_commit: st ) else: # The comment does not have code annotations - self.boto_client.post_comment_for_pull_request( + return self.boto_client.post_comment_for_pull_request( pullRequestId=str(pr_number), repositoryName=repo_name, beforeCommitId=destination_commit, @@ -277,3 +278,66 @@ def publish_comment(self, repo_name: str, pr_number: int, destination_commit: st raise ValueError("Boto3 client error calling post_comment_for_pull_request") from e except Exception as e: raise ValueError("Error calling post_comment_for_pull_request") from e + + def get_comments_for_pull_request(self, pr_number: int): + """ + Retrieve all comments for a pull request. + + Args: + - pr_number: The AWS CodeCommit pull request number + + Returns: + - The flattened commentsForPullRequestData entries from all result pages + + Boto3 Documentation: + - aws codecommit get_comments_for_pull_request + - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/codecommit/client/get_comments_for_pull_request.html + """ + if self.boto_client is None: + self._connect_boto_client() + + comments_for_pull_request = [] + try: + paginator = self.boto_client.get_paginator("get_comments_for_pull_request") + for page in paginator.paginate( + pullRequestId=str(pr_number), + PaginationConfig={"PageSize": self.comments_page_size}, + ): + comments_for_pull_request.extend(page.get("commentsForPullRequestData", [])) + except botocore.exceptions.ClientError as e: + if e.response["Error"]["Code"] == 'PullRequestDoesNotExistException': + raise ValueError(f"PR number does not exist: {pr_number}") from e + raise ValueError("Boto3 client error calling get_comments_for_pull_request") from e + except Exception as e: + raise ValueError("Error calling get_comments_for_pull_request") from e + + return comments_for_pull_request + + def update_comment(self, comment_id: str, comment: str): + """ + Update an existing pull request comment. + + Args: + - comment_id: The system-generated CodeCommit comment identifier + - comment: The replacement comment body + + Returns: + - The boto3 update_comment response + + Boto3 Documentation: + - aws codecommit update_comment + - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/codecommit/client/update_comment.html + """ + if self.boto_client is None: + self._connect_boto_client() + + try: + return self.boto_client.update_comment(commentId=str(comment_id), content=comment) + except botocore.exceptions.ClientError as e: + if e.response["Error"]["Code"] == 'CommentDoesNotExistException': + raise ValueError(f"CodeCommit comment does not exist: {comment_id}") from e + if e.response["Error"]["Code"] == 'CommentDeletedException': + raise ValueError(f"CodeCommit comment has been deleted: {comment_id}") from e + raise ValueError("Boto3 client error calling update_comment") from e + except Exception as e: + raise ValueError("Error calling update_comment") from e diff --git a/pr_agent/git_providers/codecommit_provider.py b/pr_agent/git_providers/codecommit_provider.py index 6f6b8350a9..c2878d4789 100644 --- a/pr_agent/git_providers/codecommit_provider.py +++ b/pr_agent/git_providers/codecommit_provider.py @@ -1,14 +1,22 @@ import os import re from collections import Counter +from datetime import datetime +from types import SimpleNamespace from typing import List, Optional, Tuple from urllib.parse import urlparse from pr_agent.algo.language_handler import is_valid_file +from pr_agent.algo.review_finding_state import split_review_state_marker from pr_agent.algo.types import EDIT_TYPE, FilePatchInfo from pr_agent.git_providers.codecommit_client import CodeCommitClient -from ..algo.utils import load_large_diff +from ..algo.utils import ( + add_pr_review_identity, + comment_carries_other_identity, + comment_matches_identity, + load_large_diff, +) from ..config_loader import get_settings from ..log import get_logger from .git_provider import GitProvider @@ -62,6 +70,15 @@ class CodeCommitProvider(GitProvider): This class implements the GitProvider interface for AWS CodeCommit repositories. """ + # PostCommentForPullRequest / UpdateComment reject a body above 10,240 + # characters and raise instead of degrading (#3272). Every outgoing body + # goes through _prepare_comment_body, which caps it AFTER the newline + # doubling and after any persistent-comment header has been added, so the + # cap is measured on what CodeCommit actually receives. Class-level, like + # the other providers' max_comment_length, minus the truncation marker + # limit_output_characters appends. + max_comment_length = 10240 - len("...") + def __init__(self, pr_url: Optional[str] = None, incremental: Optional[bool] = False): self.codecommit_client = CodeCommitClient() self.aws_client = None @@ -79,7 +96,6 @@ def provider_name(self): def is_supported(self, capability: str) -> bool: if capability in [ - "get_issue_comments", "create_inline_comment", "publish_inline_comments", "get_labels", @@ -195,31 +211,102 @@ def publish_description(self, pr_title: str, pr_body: str): def publish_comment(self, pr_comment: str, is_temporary: bool = False): if is_temporary: get_logger().info(pr_comment) - return - - pr_comment = CodeCommitProvider._remove_markdown_html(pr_comment) - pr_comment = CodeCommitProvider._add_additional_newlines(pr_comment) + return None try: - for target in self._get_target_contexts(): - self.codecommit_client.publish_comment( - repo_name=target["repository_name"], - pr_number=self.pr_num, - destination_commit=target["destination_commit"], - source_commit=target["source_commit"], - comment=pr_comment, - ) + published_comments = [ + self._publish_comment_to_target(pr_comment, target) + for target in self._get_target_contexts() + ] + if len(published_comments) == 1: + return published_comments[0] + return published_comments except Exception as e: raise ValueError(f"CodeCommit Cannot publish comment for PR: {self.pr_num}") from e + def publish_persistent_comment(self, pr_comment: str, + initial_header: str, + update_header: bool = True, + name='review', + final_update_message=True, + as_thread: bool = False, + identity_marker: str | None = None, + legacy_initial_header: str | None = None): + if as_thread: + get_logger().debug("CodeCommit does not support threaded persistent comments; publishing as a PR comment") + + if not identity_marker: + get_logger().debug( + "CodeCommit persistent comment updates require a stable identity marker; publishing a new comment" + ) + return self.publish_comment(pr_comment) + + persistent_comment = add_pr_review_identity(pr_comment, identity_marker) + try: + comments = self.get_issue_comments_newest_first() + except Exception as e: + get_logger().warning( + f"CodeCommit could not read existing comments; publishing a new persistent comment: {e}" + ) + return self.publish_comment(persistent_comment) + + identifiers = [identity_marker, legacy_initial_header] + used_comment_ids = set() + published_comments = [] + target_contexts = self._get_target_contexts() + destination_counts = Counter( + (target["repository_name"], target["destination_commit"]) + for target in target_contexts + ) + allow_repository_fallback = len(target_contexts) == 1 + + for target in target_contexts: + comment_to_update = self._find_persistent_comment_for_target( + comments, + target, + identifiers, + identity_marker, + used_comment_ids, + destination_counts[(target["repository_name"], target["destination_commit"])] == 1, + allow_repository_fallback, + ) + if comment_to_update is None: + published_comments.append(self._publish_comment_to_target(persistent_comment, target)) + continue + + used_comment_ids.add(comment_to_update.id) + comment_body = self._persistent_body_for_target( + persistent_comment, + update_header, + name, + identity_marker, + target, + ) + comment_url = self.get_comment_url(comment_to_update) + get_logger().info(f"Persistent mode - updating comment {comment_url} to latest {name} message") + if self.edit_comment(comment_to_update, comment_body) is not False: + published_comments.append(comment_to_update) + continue + + published_comments.append(self._publish_comment_to_target(persistent_comment, target)) + + if final_update_message: + get_logger().debug("CodeCommit does not publish separate persistent update status comments") + if len(published_comments) == 1: + return published_comments[0] + return published_comments + def publish_code_suggestions(self, code_suggestions: list) -> bool: counter = 1 + publishable_count = 0 + published_count = 0 for suggestion in code_suggestions: # Verify that each suggestion has the required keys if not all(key in suggestion for key in ["body", "relevant_file", "relevant_lines_start"]): get_logger().warning(f"Skipping code suggestion #{counter}: Each suggestion must have 'body', 'relevant_file', 'relevant_lines_start' keys") continue + publishable_count += 1 target_contexts = self._get_target_contexts_for_file(suggestion["relevant_file"]) for target in target_contexts: try: @@ -236,15 +323,15 @@ def publish_code_suggestions(self, code_suggestions: list) -> bool: annotation_file=suggestion["relevant_file"], annotation_line=suggestion["relevant_lines_start"], ) + published_count += 1 except Exception as e: raise ValueError(f"CodeCommit Cannot publish code suggestions for PR: {self.pr_num}") from e counter += 1 - # The calling function passes in a list of code suggestions, and this function publishes each suggestion one at a time. - # If we were to return False here, the calling function will attempt to publish the same list of code suggestions again, one at a time. - # Since this function publishes the suggestions one at a time anyway, we always return True here to avoid the retry. - return True + # A partial failure must not report failure: the caller republishes the whole + # list, which would post the already-accepted suggestions a second time. + return published_count > 0 or publishable_count == 0 def publish_labels(self, labels): return [""] # not implemented yet @@ -258,6 +345,26 @@ def remove_initial_comment(self): def remove_comment(self, comment): return "" # not implemented yet + def edit_comment(self, comment, body: str): + comment_id = comment.get("id") if isinstance(comment, dict) else getattr(comment, "id", None) + if not isinstance(comment_id, str) or not comment_id: + get_logger().warning(f"CodeCommit cannot update comment without a valid comment id: {comment_id!r}") + return False + + body = self._prepare_comment_body(body) + try: + response = self.codecommit_client.update_comment(comment_id, body) + except Exception as e: + get_logger().warning(f"CodeCommit failed to update comment {comment_id}: {e}") + return False + + updated_comment = response.get("comment") if isinstance(response, dict) else None + if not isinstance(updated_comment, dict): + return False + if updated_comment.get("deleted") is True: + return False + return updated_comment.get("commentId") == comment_id and updated_comment.get("content") == body + def publish_inline_comment(self, body: str, relevant_file: str, relevant_line_in_file: str, original_suggestion=None): # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/codecommit/client/post_comment_for_compared_commit.html raise NotImplementedError("CodeCommit provider does not support publishing inline comments yet") @@ -323,7 +430,19 @@ def get_user_id(self): return -1 # not implemented yet def get_issue_comments(self): - raise NotImplementedError("CodeCommit provider does not support issue comments yet") + comments = [] + for comment_data in self.codecommit_client.get_comments_for_pull_request(self.pr_num): + comments.extend(self._extract_issue_comments(comment_data)) + return sorted(comments, key=self._comment_sort_key) + + def get_issue_comments_newest_first(self): + return sorted(self.get_issue_comments(), key=self._comment_sort_key, reverse=True) + + def supports_review_comment_identity(self) -> bool: + return True + + def get_comment_url(self, comment) -> str: + return self.get_pr_url() def get_repo_settings(self): # a local ".pr_agent.toml" settings file is optional @@ -464,6 +583,178 @@ def _get_target_contexts_for_file(self, filename: str): def get_commit_messages(self) -> str: return "" # not implemented yet + def _publish_comment_to_target(self, pr_comment: str, target: dict): + pr_comment = self._prepare_comment_body(pr_comment) + response = self.codecommit_client.publish_comment( + repo_name=target["repository_name"], + pr_number=self.pr_num, + destination_commit=target["destination_commit"], + source_commit=target["source_commit"], + comment=pr_comment, + ) + if isinstance(response, dict): + return self._comment_from_api_comment( + response.get("comment"), + { + "repositoryName": target["repository_name"], + "beforeCommitId": target["destination_commit"], + "afterCommitId": target["source_commit"], + }, + ) + return response + + def _find_persistent_comment_for_target( + self, + comments: list, + target: dict, + identifiers: list, + identity_marker: str, + used_comment_ids: set, + allow_destination_fallback: bool, + allow_repository_fallback: bool, + ): + matchers = [self._comment_matches_target_exact] + if allow_destination_fallback: + matchers.append(self._comment_matches_target_destination) + if allow_repository_fallback: + matchers.append(self._comment_matches_target_repository) + + for matcher in matchers: + comment = self._find_persistent_comment( + comments, + identifiers, + identity_marker, + used_comment_ids, + lambda candidate: matcher(candidate, target), + ) + if comment is not None: + return comment + return None + + @staticmethod + def _find_persistent_comment( + comments: list, + identifiers: list, + identity_marker: str, + used_comment_ids: set, + target_matcher, + ): + for identifier in identifiers: + if not identifier: + continue + for comment in comments: + if comment.id in used_comment_ids or not target_matcher(comment): + continue + body = GitProvider._get_comment_body(comment) + if not comment_matches_identity(body, identifier): + continue + if comment_carries_other_identity(body, identity_marker): + continue + return comment + return None + + @staticmethod + def _comment_matches_target_exact(comment, target: dict) -> bool: + return ( + CodeCommitProvider._comment_matches_target_destination(comment, target) + and getattr(comment, "after_commit_id", None) == target["source_commit"] + ) + + @staticmethod + def _comment_matches_target_destination(comment, target: dict) -> bool: + return ( + CodeCommitProvider._comment_matches_target_repository(comment, target) + and getattr(comment, "before_commit_id", None) == target["destination_commit"] + ) + + @staticmethod + def _comment_matches_target_repository(comment, target: dict) -> bool: + repository_name = getattr(comment, "repository_name", None) + return repository_name == target["repository_name"] + + @staticmethod + def _persistent_body_for_target( + pr_comment: str, + update_header: bool, + name: str, + identity_marker: str, + target: dict, + ) -> str: + if not update_header: + return pr_comment + + update_message = f"#### ({name.capitalize()} updated until commit {target['source_commit']})\n" + updated_anchor = f"{identity_marker}\n\n{update_message}" + return pr_comment.replace(identity_marker, updated_anchor, 1) + + def _prepare_comment_body(self, pr_comment: str) -> str: + pr_comment = CodeCommitProvider._remove_markdown_html(pr_comment) + body, marker = split_review_state_marker(pr_comment) + body = CodeCommitProvider._add_additional_newlines(body) + if not marker: + return self.limit_output_characters(body, self.max_comment_length) + # The persistent review state is a hidden marker at the end of the body. + # The reviewer sizes it before the newline doubling above, so the doubled + # body can overrun the cap; truncate only the human text and keep the + # marker whole, otherwise the next run cannot parse the state. + budget = self.max_comment_length - len(marker) - 2 + if budget <= 0: + return marker[: self.max_comment_length] + return f"{self.limit_output_characters(body, budget)}\n\n{marker}" + + @staticmethod + def _extract_issue_comments(comment_data: dict): + if not isinstance(comment_data, dict) or comment_data.get("location"): + return [] + + comments = [] + for comment in comment_data.get("comments", []): + issue_comment = CodeCommitProvider._comment_from_api_comment(comment, comment_data) + if issue_comment is not None: + comments.append(issue_comment) + return comments + + @staticmethod + def _comment_from_api_comment(comment: dict, comment_data: dict): + if not isinstance(comment, dict): + return None + if comment.get("deleted") is True or comment.get("inReplyTo"): + return None + + body = comment.get("content") + comment_id = comment.get("commentId") + created_at = CodeCommitProvider._comment_timestamp(comment.get("creationDate")) + if not isinstance(body, str) or not body: + return None + if not isinstance(comment_id, str) or not comment_id: + return None + if created_at is None: + return None + + author_arn = comment.get("authorArn") if isinstance(comment.get("authorArn"), str) else "" + return SimpleNamespace( + body=body, + id=comment_id, + created_at=created_at, + last_modified_at=CodeCommitProvider._comment_timestamp(comment.get("lastModifiedDate")), + repository_name=comment_data.get("repositoryName"), + before_commit_id=comment_data.get("beforeCommitId"), + after_commit_id=comment_data.get("afterCommitId"), + user=SimpleNamespace(login=author_arn), + ) + + @staticmethod + def _comment_timestamp(value): + if isinstance(value, datetime): + return value.timestamp() + if isinstance(value, (int, float)): + return float(value) + return None + + @staticmethod + def _comment_sort_key(comment): + return (comment.created_at, comment.id) + @staticmethod def _add_additional_newlines(body: str) -> str: """ diff --git a/pr_agent/git_providers/gerrit_provider.py b/pr_agent/git_providers/gerrit_provider.py index aa5603b770..dd9d0e4b1b 100644 --- a/pr_agent/git_providers/gerrit_provider.py +++ b/pr_agent/git_providers/gerrit_provider.py @@ -387,8 +387,10 @@ def split_suggestion(self, msg) -> tuple[str, str]: '\n'.join(context) + '\n' if context else '' ) - def publish_code_suggestions(self, code_suggestions: list): + def publish_code_suggestions(self, code_suggestions: list) -> bool: msg = [] + publishable_count = 0 + published_count = 0 repo_root = pathlib.Path(self.repo_path).resolve() for suggestion in code_suggestions: # Validate suggestion structure before accessing keys @@ -402,6 +404,8 @@ def publish_code_suggestions(self, code_suggestions: list): except ValueError: get_logger().warning(f"Skipping suggestion with path traversal: {suggestion['relevant_file']}") continue + + publishable_count += 1 description, code = self.split_suggestion(suggestion['body']) add_suggestion( target_path, @@ -417,8 +421,13 @@ def publish_code_suggestions(self, code_suggestions: list): msg.append(f'* {description}\n{full_path}') if msg: - add_comment(self.parsed_url, self.refspec, "\n".join(msg)) - return True + try: + add_comment(self.parsed_url, self.refspec, "\n".join(msg)) + published_count += 1 + except Exception as e: + get_logger().exception("Failed to publish Gerrit code suggestions: {}", e) + + return published_count > 0 or publishable_count == 0 def publish_comment(self, pr_comment: str, is_temporary: bool = False): if not is_temporary: diff --git a/pr_agent/git_providers/git_provider.py b/pr_agent/git_providers/git_provider.py index 7c11bf8de5..00ab7f8b1f 100644 --- a/pr_agent/git_providers/git_provider.py +++ b/pr_agent/git_providers/git_provider.py @@ -1,4 +1,3 @@ -# enum EDIT_TYPE (ADDED, DELETED, MODIFIED, RENAMED) import os import re import shutil @@ -8,6 +7,7 @@ from collections.abc import Iterable from typing import Optional, Tuple +from pr_agent.algo.language_handler import numeric_languages from pr_agent.algo.types import FilePatchInfo from pr_agent.algo.utils import ( Range, @@ -184,6 +184,67 @@ def supports_threaded_pr_questions(self) -> bool: def supports_line_question_history(self) -> bool: return False + def supports_checkbox_commands(self) -> bool: + """Whether a published comment renders command checkboxes as checkboxes. + + Providers that render `- [ ]` as a tickable box override this; the default is no + support, so tools render commands as text instead.""" + return False + + def supports_pr_chat(self) -> bool: + """Whether this provider is compatible with the linked PR-Agent browser-extension chat experience.""" + return False + + @classmethod + def supports_issue_indexing(cls) -> bool: + """Whether `/similar_issue` can read and index this provider's issues. + + Declared on the class rather than on an instance because the tool consults it before + constructing a provider: `PRSimilarIssue` needs to know whether to build one at all. + The indexing path relies on issue listing, issue bodies and issue comments, so a + provider that exposes those overrides this; the default is no support, so the tool + reports the command as unsupported instead of failing part-way through. + """ + return False + + def supports_inline_help_footer(self) -> bool: + """Whether the `/describe` help footer is rendered as an inline `
  • ` list. + + Scoped to that footer's layout, not to HTML lists in general: the changes + walkthrough already emits `
      ` and `
    • ` for every provider that passes the + `gfm_markdown` gate. Providers whose `
      ` summary renders a sibling `
    • ` + inline override this; the default falls back to `
      `-separated bullets.""" + return False + + def supports_changelog_update_review(self) -> bool: + """Whether a pushed CHANGELOG.md commit can be annotated with a PR review. + + `/update_changelog --push_changelog_changes=true` posts its summary as a review on the + commit it just pushed. Providers exposing a commit-scoped review API override this; + the default is no support, so the review is simply skipped. + """ + return False + + def supports_markdown_tables(self) -> bool: + """Whether comments render pipe-table markdown. + + Only consulted for providers without `gfm_markdown`, so that tools can degrade to + a plain table instead of refusing to render. Providers that render Markdown tables + but not GitHub-flavored markdown override this.""" + return False + + def supports_issue_url_tickets(self) -> bool: + """Tickets are linked as issue URLs in the PR description or branch name.""" + return False + + def supports_issue_reference_tickets(self) -> bool: + """Tickets are linked as project-scoped issue references (e.g. group/project#12).""" + return False + + def supports_linked_work_item_tickets(self) -> bool: + """Tickets come from work items the platform links to the PR itself.""" + return False + #Given a url (issues or PR/MR) - get the .git repo url to which they belong. Needs to be implemented by the provider. def get_git_repo_url(self, issues_or_pr_url: str) -> str: get_logger().warning("Not implemented! Returning empty url") @@ -217,13 +278,9 @@ def _prepare_clone_url_with_token(self, repo_url_to_clone: str) -> str | None: get_logger().warning("Not implemented! Returning None") return None - # Does a shallow clone, using a forked process to support a timeout guard. - # In case operation has failed, it is expected to throw an exception as this method does not return a value. + # Run a shallow, blob-filtered clone in a subprocess so the timeout can terminate the operation. + # Failures propagate to clone(), which handles and logs them. def _clone_inner(self, repo_url: str, dest_folder: str, operation_timeout_in_seconds: int=None) -> None: - #The following ought to be equivalent to: - # #Repo.clone_from(repo_url, dest_folder) - # , but with throwing an exception upon timeout. - # Note: This can only be used in context that supports using pipes. try: ssl_env = get_git_ssl_env() except Exception as e: @@ -266,6 +323,17 @@ def clone(self, repo_url_to_clone: str, dest_folder: str, remove_dest_folder: bo def get_files(self) -> list: pass + def get_pr_file_paths(self) -> list: + """Return every repository-relative path the PR/MR touches, independent of + incremental review state, preserving rename metadata. + + The default delegates to get_files(). Providers whose get_files() shrinks + to the unreviewed subset while an incremental review is active must + override this with a complete file-set listing, so per-directory settings + discovery does not depend on how much of the PR the review has covered. + """ + return self.get_files() + @abstractmethod def get_diff_files(self) -> list[FilePatchInfo]: pass @@ -416,6 +484,21 @@ def _is_generated_by_pr_agent(self, description_lowercase: str) -> bool: def get_repo_settings(self): pass + def get_repo_settings_tree(self, ref: str = "") -> tuple[list[str], str]: + """Recursively list every `.pr_agent.toml` path at `ref` ("" = the repository + default branch) as `(paths, resolved_ref)`. Providers without per-directory + settings support return `([], "")` so the feature degrades to root-only + behavior. Implemented by GitHub and GitLab.""" + return [], "" + + def get_repo_settings_contents(self, paths: list[str], ref: str) -> dict[str, bytes]: + """Fetch the raw content of per-directory repo settings files at `ref`. + + Only the entries whose content was fetched successfully are returned; a + missing file is skipped with a warning rather than failing the request. + Defaults to no per-directory support.""" + return {} + def get_owning_namespace(self) -> Optional[str]: """Return the org/group/workspace that owns this repository, or None when the provider has no organisation-level home for global settings. @@ -463,6 +546,51 @@ def _fetch_global_repo_settings(self, namespace: str): def get_repo_file_content(self, file_path: str, from_default_branch: bool = False): return "" + def get_sibling_repo_file_content(self, repo_id: str, file_path: str, from_default_branch: bool = False): + """Fetch a single file from a sibling repository in the same namespace/owner. + + Used by repo context when a repo_context_files entry is a + sibling dict ``{"repo_id": ..., "file_path": ...}``. Only providers that can resolve + the sibling through their own authenticated API (GitHub, GitLab) override this; both + require host allowlisting, check the resolved owner/group, and read from the + sibling's default branch. The default returns "" so unsupported providers degrade + gracefully without reaching an unrelated repository or host. + """ + return "" + + def is_sibling_repo_allowed(self, repo_id: str, *, case_sensitive: bool = True) -> bool: + """Require explicit host approval before resolving a sibling repository.""" + allowed = get_settings().config.get("repo_context_sibling_repos", []) + if not isinstance(allowed, list): + return False + normalize = (lambda value: value) if case_sensitive else str.casefold + return normalize(repo_id) in { + normalize(value.strip().strip("/")) for value in allowed if isinstance(value, str) and value.strip() + } + + def set_command_actor(self, actor) -> None: + """Record the authenticated user who triggered the current command. + + Comment commands can pass arbitrary arguments, so sibling-repo context must be + authorized against the actor who issued the command rather than the PR/MR author: + a commenter may not have the read access the author has. Providers that resolve + siblings through their own authenticated API use this identity (when set) instead + of the PR/MR author. When no trustworthy actor is available the providers fail + closed for non-public siblings. + """ + self._command_actor = actor + + def get_repo_context_ref(self, from_default_branch: bool = False) -> Optional[str]: + """Return the ref (commit SHA or branch name) that repo-context files are read from. + + The repo-context cache key (pr_agent/algo/repo_context.py) includes this ref so a + rebase or a push to the base branch invalidates cached file content instead of serving + it from a commit that has since moved. Providers that override get_repo_file_content + should return the same ref they fetch from; the default None covers providers with no + repo-context support at all. + """ + return None + def get_workspace_name(self): return "" @@ -521,6 +649,16 @@ def supports_thread_resolution(self) -> bool: def resolve_outdated_inline_threads(self): # noqa: B027 - intentional no-op pass + def supports_comment_editing(self) -> bool: + """Whether this provider can actually edit an existing comment. + + The base ``edit_comment`` is a no-op that returns ``None``, which + ``publish_persistent_comment_full`` cannot distinguish from a successful + edit. A provider that has not implemented it therefore cannot persist, + and must create a new comment instead of silently discarding the body. + """ + return type(self).edit_comment is not GitProvider.edit_comment + def publish_persistent_comment(self, pr_comment: str, initial_header: str, update_header: bool = True, @@ -529,7 +667,18 @@ def publish_persistent_comment(self, pr_comment: str, as_thread: bool = False, identity_marker: str | None = None, legacy_initial_header: str | None = None): - return self.publish_comment(pr_comment, **({'as_thread': True} if as_thread else {})) + if not self.supports_comment_editing(): + return self.publish_comment(pr_comment, **({'as_thread': True} if as_thread else {})) + return self.publish_persistent_comment_full( + pr_comment, + initial_header, + update_header, + name, + final_update_message, + as_thread=as_thread, + identity_marker=identity_marker, + legacy_initial_header=legacy_initial_header, + ) @staticmethod def _get_comment_body(comment) -> str: @@ -754,7 +903,14 @@ def get_num_of_files(self): return -1 def limit_output_characters(self, output: str, max_chars: int): - return output[:max_chars] + '...' if len(output) > max_chars else output + """Truncate output to max_chars, including the truncation suffix.""" + if len(output) <= max_chars: + return output + if max_chars <= 0: + return "" + suffix = "..." + suffix = suffix[:max_chars] + return output[:max_chars - len(suffix)] + suffix def get_main_pr_language(languages, files) -> str: @@ -770,6 +926,9 @@ def get_main_pr_language(languages, files) -> str: return main_language_str try: + languages = numeric_languages(languages) + if not languages: + return main_language_str top_language = max(languages, key=languages.get).lower() # validate that the specific commit uses the main language diff --git a/pr_agent/git_providers/gitea_provider.py b/pr_agent/git_providers/gitea_provider.py index 27e9d793fa..1d688b90d8 100644 --- a/pr_agent/git_providers/gitea_provider.py +++ b/pr_agent/git_providers/gitea_provider.py @@ -322,24 +322,6 @@ def get_comment_url(self, comment) -> str: return comment.get("html_url") or comment.get("url") or "" return getattr(comment, "html_url", "") or getattr(comment, "url", "") - def publish_persistent_comment(self, pr_comment: str, - initial_header: str, - update_header: bool = True, - name='review', - final_update_message=True, - identity_marker: str | None = None, - legacy_initial_header: str | None = None): - # Keep the legacy updater path until Gitea normalizes its dictionary-shaped comment payloads. - return self.publish_persistent_comment_full( - pr_comment, - initial_header, - update_header, - name, - final_update_message, - identity_marker=identity_marker, - legacy_initial_header=legacy_initial_header, - ) - def publish_comment(self, comment: str,is_temporary: bool = False) -> None: """Publish a comment to the pull request""" if is_temporary and not get_settings().config.publish_output_progress: @@ -953,6 +935,12 @@ def get_repo_file_content(self, file_path: str, from_default_branch: bool = Fals return "" raise + def get_repo_context_ref(self, from_default_branch: bool = False) -> Optional[str]: + if from_default_branch: + return self.repo_api.repo_get(self.owner, self.repo).default_branch + # Only trust the PR target (base) ref — never the PR head (self.sha). + return self.base_sha or self.base_ref + class RepoApi(giteapy.RepositoryApi): def __init__(self, client: giteapy.ApiClient): self.repository = giteapy.RepositoryApi(client) diff --git a/pr_agent/git_providers/github_provider.py b/pr_agent/git_providers/github_provider.py index eb9eb54d09..1aa82abdb2 100644 --- a/pr_agent/git_providers/github_provider.py +++ b/pr_agent/git_providers/github_provider.py @@ -11,7 +11,7 @@ from typing import Optional, Tuple from urllib.parse import urlparse -from github import AppAuthentication, Auth, Github, GithubException, GithubIntegration +from github import Auth, Github, GithubException, GithubIntegration, GithubRetry from github.Issue import Issue from retry.api import retry_call from starlette_context import context @@ -78,6 +78,7 @@ def __init__(self, pr_url: Optional[str] = None): self.diff_files = None self.git_files = None self.incremental = IncrementalPR(False) + self._resolved_config_branch: str | None = None self._check_run_ids: dict = {} if pr_url and 'pull' in pr_url: self.set_pr(pr_url) @@ -125,6 +126,25 @@ def is_supported(self, capability: str) -> bool: def supports_line_question_history(self) -> bool: return True + def supports_checkbox_commands(self) -> bool: + return True + + def supports_inline_help_footer(self) -> bool: + return True + + def supports_pr_chat(self) -> bool: + return True + + @classmethod + def supports_issue_indexing(cls) -> bool: + return True + + def supports_changelog_update_review(self) -> bool: + return True + + def supports_issue_url_tickets(self) -> bool: + return True + def _get_owner_and_repo_path(self, given_url: str) -> str: try: repo_path = None @@ -245,6 +265,30 @@ def get_files(self): self.git_files = list(self.pr.get_files()) return self.git_files + def get_pr_file_paths(self): + """Return the complete PR file set regardless of incremental review state. + + get_files() returns only the unreviewed subset once an incremental review + is active, so per-directory settings would change between commands based on + which files the review already covered. Discovery instead walks the full PR + file set, preserving rename metadata (previous_filename so both sides of a + move apply). Reuses the same context["git_files"] cache as get_files() and + never falls back to the incremental-aware listing. + """ + try: + git_files = context.get("git_files", None) + if git_files: + return git_files + if getattr(self, "git_files", None): + return self.git_files + git_files = list(self.pr.get_files()) + context["git_files"] = git_files + return git_files + except Exception: + if getattr(self, "git_files", None): + return self.git_files + return list(self.pr.get_files()) + def get_num_of_files(self): if hasattr(self.git_files, "totalCount"): return self.git_files.totalCount @@ -414,17 +458,19 @@ def publish_persistent_comment(self, pr_comment: str, update_header: bool = True, name='review', final_update_message=True, + as_thread: bool = False, identity_marker: str | None = None, legacy_initial_header: str | None = None): if get_settings().github.publish_as_check_run: if self._publish_check_run(pr_comment, name): return - self.publish_persistent_comment_full( + return self.publish_persistent_comment_full( pr_comment, initial_header, update_header, name, final_update_message, + as_thread=as_thread, identity_marker=identity_marker, legacy_initial_header=legacy_initial_header, ) @@ -460,8 +506,10 @@ def _resolve_app_login(self) -> str: return cached try: integration = GithubIntegration( - integration_id=str(get_settings().github.app_id), - private_key=get_settings().github.private_key, + auth=Auth.AppAuth( + app_id=str(get_settings().github.app_id), + private_key=get_settings().github.private_key, + ), base_url=self.base_url, ) slug = (getattr(integration.get_app(), "slug", "") or "").strip() @@ -694,7 +742,7 @@ def publish_inline_comments(self, comments: list[dict], disable_fallback: bool = get_logger().info( f"Persistent inline comments: all {skipped} suggestion(s) " f"already posted; nothing to publish") - return + return True comments = deduped else: comments = [ @@ -713,6 +761,7 @@ def publish_inline_comments(self, comments: list[dict], disable_fallback: bool = for body_fp, code_fp in pending_fingerprints: store.add(body_fp) store.add(code_fp) + return True except Exception as e: get_logger().info("Initially failed to publish inline comments as committable") @@ -722,7 +771,8 @@ def publish_inline_comments(self, comments: list[dict], disable_fallback: bool = raise e # will end up with publishing the comments one by one try: - self._publish_inline_comments_fallback_with_verification(comments) + published_count = self._publish_inline_comments_fallback_with_verification(comments) + return bool(published_count) except Exception as e: get_logger().error(f"Failed to publish inline code comments fallback, error: {e}") raise @@ -887,11 +937,13 @@ def _publish_inline_comments_fallback_with_verification(self, comments: list[dic then publish all the remaining valid comments in a single review. For invalid comments, also try removing the suggestion part and posting the comment just on the first line. """ + published_count = 0 verified_comments, invalid_comments = self._verify_code_comments(comments) # publish as a group the verified comments if verified_comments: self.pr.create_review(commit=self.last_commit_id, comments=verified_comments) + published_count += len(verified_comments) # try to publish one by one the invalid comments as a one-line code comment if invalid_comments and get_settings().github.try_fix_invalid_inline_comments: @@ -899,9 +951,10 @@ def _publish_inline_comments_fallback_with_verification(self, comments: list[dic fixed_comments_as_one_liner = self._try_fix_invalid_inline_comments(invalid_comments_list) for comment in fixed_comments_as_one_liner: try: - self.publish_inline_comments([comment], disable_fallback=True) - get_logger().info(f"Published invalid comment as a single line comment: {comment}") - except: + if self.publish_inline_comments([comment], disable_fallback=True): + published_count += 1 + get_logger().info(f"Published invalid comment as a single line comment: {comment}") + except Exception: get_logger().error(f"Failed to publish invalid comment as a single line comment: {comment}") dropped_count = len(invalid_comments) - len(fixed_comments_as_one_liner) @@ -920,6 +973,7 @@ def _publish_inline_comments_fallback_with_verification(self, comments: list[dic f"Dropped {len(invalid_comments)} invalid comments " f"(try_fix_invalid_inline_comments is off). Paths: {dropped_paths}" ) + return published_count def _verify_code_comment(self, comment: dict): is_verified = False @@ -1029,8 +1083,7 @@ def publish_code_suggestions(self, code_suggestions: list) -> bool: post_parameters_list.append(post_parameters) try: - self.publish_inline_comments(post_parameters_list) - return True + return bool(self.publish_inline_comments(post_parameters_list)) except Exception as e: get_logger().error(f"Failed to publish code suggestion, error: {e}") return False @@ -1085,40 +1138,6 @@ def get_comment_body_from_comment_id(self, comment_id: int): get_logger().exception(f"Failed to edit comment, error: {e}") return None - def publish_file_comments(self, file_comments: list) -> bool: - try: - headers, existing_comments = self.pr._requester.requestJsonAndCheck( - "GET", f"{self.pr.url}/comments" - ) - for comment in file_comments: - comment['commit_id'] = self.last_commit_id.sha - comment['body'] = self.limit_output_characters(comment['body'], self.max_comment_chars) - - found = False - for existing_comment in existing_comments: - comment['commit_id'] = self.last_commit_id.sha - our_app_name = get_settings().get("GITHUB.APP_NAME", "") - same_comment_creator = False - if self.deployment_type == 'app': - same_comment_creator = our_app_name.lower() in existing_comment['user']['login'].lower() - elif self.deployment_type == 'user': - same_comment_creator = self.github_user_id == existing_comment['user']['login'] - if existing_comment['subject_type'] == 'file' and comment['path'] == existing_comment['path'] and same_comment_creator: - - headers, data_patch = self.pr._requester.requestJsonAndCheck( - "PATCH", f"{self.base_url}/repos/{self.repo}/pulls/comments/{existing_comment['id']}", input={"body":comment['body']} - ) - found = True - break - if not found: - headers, data_post = self.pr._requester.requestJsonAndCheck( - "POST", f"{self.pr.url}/comments", input=comment - ) - return True - except Exception as e: - get_logger().error(f"Failed to publish diffview file summary, error: {e}") - return False - def remove_initial_comment(self): try: for comment in getattr(self.pr, 'comments_list', []): @@ -1148,12 +1167,13 @@ def get_pr_owner_id(self) -> str | None: return None return self.repo.split('/')[0] - def get_owning_namespace(self) -> Optional[str]: + def get_owning_namespace(self, *, resolved: bool = False) -> Optional[str]: # Be robust to providers built without full __init__ (e.g. __new__ in tests/helpers): # without a repo there is no org to resolve, so skip global settings quietly. if not getattr(self, "repo", None): return None - return self.repo.split('/')[0] + repo_path = self.github_client.get_repo(self.repo).full_name if resolved else self.repo + return repo_path.split("/")[0] if isinstance(repo_path, str) and "/" in repo_path else None def get_pr_description_full(self): return self.pr.body @@ -1197,6 +1217,7 @@ def get_repo_settings(self): # left to propagate so they aren't masked by a silent fallback. try: contents = self.repo_obj.get_contents(".pr_agent.toml", ref=config_branch).decoded_content + self._resolved_config_branch = config_branch if settings_files: settings_files.append(("local", contents)) return settings_files @@ -1212,6 +1233,7 @@ def get_repo_settings(self): try: # more logical to take 'pr_agent.toml' from the default branch contents = self.repo_obj.get_contents(".pr_agent.toml").decoded_content + self._resolved_config_branch = getattr(self.repo_obj, "default_branch", "") or "" if config_branch and not settings_files: return contents settings_files.append(("local", contents)) @@ -1227,6 +1249,89 @@ def get_repo_settings(self): return settings_files if settings_files else "" + def get_repo_settings_tree(self, ref: str = "") -> tuple[list[str], str]: + """Recursively list every `.pr_agent.toml` at *ref* ("" = default branch). + + Follows the same branch resolution as get_repo_settings(): when the root + lookup resolved a config, the tree is read from that same branch + (``_resolved_config_branch``). When the root lookup resolved nothing (no + root ``.pr_agent.toml`` anywhere), the tree is read from *ref* -- the + CONFIG.CONFIG_BRANCH / PR_AGENT_CONFIG_BRANCH hint, or the repository + default branch when *ref* is empty -- falling back to the default branch + on a 404. Returns ``(paths, resolved_ref)`` where *resolved_ref* is the + branch the tree was actually read from; an empty *paths* list means the + recursive tree hit GitHub's truncation cap and per-directory settings had + to be skipped. + """ + repo = getattr(self, "repo_obj", None) + if repo is None: + return [], "" + resolved_ref = self._resolved_config_branch or ref or "" + try: + if not resolved_ref: + resolved_ref = repo.default_branch + return self._list_config_tree_paths(repo, resolved_ref), resolved_ref + except GithubException as e: + if e.status == 404 and resolved_ref: + # Branch or tree not found (possibly deleted between root and per-dir + # resolution). Fall back to the default branch; matches the root config + # fallback when CONFIG_BRANCH is stale. + get_logger().debug( + f"No git tree for branch '{resolved_ref}' while listing per-directory " + "settings; falling back to default branch" + ) + resolved_ref = repo.default_branch + return self._list_config_tree_paths(repo, resolved_ref), resolved_ref + # Unlike 404, a 403/5xx is not an expected fallback signal: propagate so it + # is not masked by a silent downgrade (matches get_repo_settings()). + raise + + def _list_config_tree_paths(self, repo, ref: str) -> list[str]: + """Fetch a recursive tree at *ref* and return its `.pr_agent.toml` blob paths. + + A truncated tree (GitHub caps recursive trees at 100k entries / 7 MB and sets + ``truncated``) cannot be trusted to name every config, so it degrades to no + per-directory settings with a warning rather than applying an incomplete, silent + subset. The root config is unaffected. + """ + tree = repo.get_git_tree(ref, recursive=True) + if getattr(tree, "truncated", False): + get_logger().warning( + f"Git tree for branch '{ref}' is truncated by GitHub's recursive-tree " + "limit; skipping per-directory settings for this repository" + ) + return [] + return self._extract_config_tree_paths(tree) + + @staticmethod + def _extract_config_tree_paths(tree) -> list[str]: + """Return repository-relative paths of every `.pr_agent.toml` blob in a + PyGithub GitTree object.""" + return [ + item.path + for item in getattr(tree, "tree", []) + if getattr(item, "type", None) == "blob" + and getattr(item, "path", "").endswith(".pr_agent.toml") + ] + + def get_repo_settings_contents(self, paths: list[str], ref: str) -> dict[str, bytes]: + """Fetch raw content of per-directory settings files at *ref*.""" + repo = getattr(self, "repo_obj", None) + if repo is None: + return {} + result: dict[str, bytes] = {} + for path in paths: + try: + result[path] = repo.get_contents(path, ref=ref).decoded_content + except GithubException as e: + if e.status == 404: + get_logger().warning( + f"Per-directory settings file '{path}' not found at ref '{ref}'; skipping" + ) + else: + raise + return result + def _get_global_settings_cache_key(self, repo_owner: str) -> str: # Cache per org AND host: the same org name on two different hosts (github.com vs a # self-hosted GitHub Enterprise instance) must not share a settings entry. @@ -1272,6 +1377,104 @@ def get_repo_file_content(self, file_path: str, from_default_branch: bool = Fals return "" raise + def get_sibling_repo_file_content(self, repo_id: str, file_path: str, from_default_branch: bool = False): + try: + repo_id = (repo_id or "").strip().strip("/") + file_path = (file_path or "").strip().lstrip("/") + if not repo_id or not file_path: + return "" + if not self.is_sibling_repo_allowed(repo_id, case_sensitive=False): + get_logger().warning(f"Ignoring sibling repo absent from the host allowlist: {repo_id}") + return "" + sibling_repo = self.github_client.get_repo(repo_id) + resolved_name = sibling_repo.full_name + current_owner = self.get_owning_namespace(resolved=True) + # Reject redirects/transfers unless the canonical repository was explicitly selected. + if (not isinstance(resolved_name, str) or resolved_name.casefold() != repo_id.casefold() + or not current_owner or resolved_name.split("/")[0].casefold() != current_owner.casefold()): + get_logger().warning(f"Ignoring out-of-owner sibling repo in repo context: {repo_id}") + return "" + if not self._requester_can_read_sibling_repo(sibling_repo): + get_logger().warning( + f"Ignoring sibling repo context file the review requester cannot read: {repo_id}" + ) + return "" + # The sibling has no PR-target ref in this repo, so its default branch is the only + # well-defined revision to read the file from. + contents = sibling_repo.get_contents(file_path).decoded_content + if isinstance(contents, bytes): + return contents.decode("utf-8", errors="replace") + return contents + except GithubException as e: + # A missing optional file is an expected "no context" outcome; transient errors + # propagate so repo context treats them as a fetch error and does not cache empties. + if e.status == 404: + return "" + raise + + def _requester_can_read_sibling_repo(self, sibling_repo) -> bool: + # Only repositories any review requester can read are granted unconditionally: a repo + # that reports no visibility and is not flagged private (i.e. public). Internal repos + # (GitHub Enterprise) are not ``private`` but are restricted to org members (and the + # outside collaborators they add), so they must be verified instead of treated as public. + visibility = getattr(sibling_repo, "visibility", None) + is_private = bool(getattr(sibling_repo, "private", False)) + if visibility == "internal": + is_private = True + if not is_private: + return True + # Private or internal: the requester must have read access. Prefer the authenticated + # command actor when one is known; otherwise (CLI runs) fall back to the PR author as + # the operator proxy, and fail closed when neither is available. + requester_login = getattr(self, "_command_actor", None) + if not requester_login: + pr = getattr(self, "pr", None) + user = getattr(pr, "user", None) if pr is not None else None + requester_login = user.get("login") if isinstance(user, dict) else getattr(user, "login", None) + if not requester_login: + return False + if getattr(getattr(sibling_repo, "owner", None), "login", None) == requester_login: + return True + if visibility == "internal": + # Every member of the owning organization can read an internal repository without a + # per-repo grant. A definitive "not a member" is *not* denial here: outside + # collaborators can be granted access to internal repositories, so fall through to + # the collaborator check on a False answer. + organization = getattr(sibling_repo, "organization", None) + if organization is not None: + try: + if bool(organization.has_in_members(self.github_client.get_user(requester_login))): + return True + except GithubException as e: + # A 404 means the organization itself cannot be resolved, so it is not a + # verdict about the requester; keep the fall-through. Auth/platform failures + # are not denials either and must surface as a fetch error, not a silent skip. + if e.status != 404: + raise + # has_in_collaborators() answers False for a definitive non-collaborator and raises for + # auth/platform failures; both are authoritative here, so transient errors propagate and + # repo context records a fetch error instead of silently dropping the sibling file. + return bool(sibling_repo.has_in_collaborators(requester_login)) + + def get_repo_context_ref(self, from_default_branch: bool = False) -> Optional[str]: + # Match get_repo_file_content: the PR target (base) commit is the cached revision. + # When the default branch is read (explicitly, or because no PR base exists) resolve + # its head commit so a push to the default branch invalidates cached content within + # the TTL instead of serving it from a moved commit. + if not from_default_branch: + base = getattr(getattr(self, "pr", None), "base", None) + ref = getattr(base, "sha", None) or getattr(base, "ref", None) + if ref: + return ref + repo_obj = getattr(self, "repo_obj", None) + if repo_obj is None: + return None + try: + return repo_obj.get_branch(repo_obj.default_branch).commit.sha + except Exception as e: + get_logger().debug(f"Could not resolve the default branch revision for repo context: {e}") + return None + def get_workspace_name(self): return self.repo.split('/')[0] @@ -1368,15 +1571,18 @@ def _get_github_client(self): if self.deployment_type == 'app': try: private_key = get_settings().github.private_key - # The app id is an integer in the settings toml, but PyJWT >=2.11 requires a - # string `iss` claim, and PyGithub 1.59 passes it through raw (#2955). + # The app id is an integer in the settings toml. PyJWT >=2.11 requires a + # string `iss` claim; PyGithub 2.7+ normalizes an int app id to a string + # upstream (#2955, PyGithub#3272), so the cast is harmless on the 2.10 pin. app_id = str(get_settings().github.app_id) except AttributeError as e: raise ValueError("GitHub app ID and private key are required when using GitHub app deployment") from e if not self.installation_id: raise ValueError("GitHub app installation ID is required when using GitHub app deployment") - auth = AppAuthentication(app_id=app_id, private_key=private_key, - installation_id=self.installation_id) + auth = Auth.AppInstallationAuth( + Auth.AppAuth(app_id=app_id, private_key=private_key), + installation_id=self.installation_id, + ) self.auth = auth elif self.deployment_type == 'user': try: @@ -1387,7 +1593,21 @@ def _get_github_client(self): "https://github.com/Codium-ai/pr-agent#method-2-run-from-source") from e self.auth = Auth.Token(token) if self.auth: - return Github(auth=self.auth, base_url=self.base_url) + github_config = get_settings().github + # PyGithub 2.x defaults to pacing and retries (0.25s between requests, 1s between + # writes, 10 retries); these had no equivalent on 1.59. The settings mirror the + # 1.59 behaviour, so the upgrade stays behaviour-neutral unless an operator opts in. + seconds_between_requests = github_config.get("seconds_between_requests", 0) + seconds_between_writes = github_config.get("seconds_between_writes", 0) + api_retries = github_config.get("api_retries", 0) + retry = GithubRetry(total=api_retries) if api_retries else None + return Github( + auth=self.auth, + base_url=self.base_url, + seconds_between_requests=seconds_between_requests, + seconds_between_writes=seconds_between_writes, + retry=retry, + ) else: raise ValueError("Could not authenticate to GitHub") @@ -1676,7 +1896,7 @@ def validate_comments_inside_hunks(self, code_suggestions): if not hasattr(file, 'patches_range'): file.patches_range = [] patch_lines = patch_str.splitlines() - for i, line in enumerate(patch_lines): + for line in patch_lines: if line.startswith('@@'): match = RE_HUNK_HEADER.match(line) # identify hunk header @@ -1696,7 +1916,7 @@ def validate_comments_inside_hunks(self, code_suggestions): min_distance = float('inf') patch_range_min = None # find the hunk that contains the comment, or the closest one - for i, patch_range in enumerate(patches_range): + for patch_range in patches_range: d1 = comment_start_line - patch_range['start'] d2 = patch_range['end'] - comment_end_line if d1 >= 0 and d2 >= 0: # found a valid hunk diff --git a/pr_agent/git_providers/gitlab_provider.py b/pr_agent/git_providers/gitlab_provider.py index 33553d880f..e8208de700 100644 --- a/pr_agent/git_providers/gitlab_provider.py +++ b/pr_agent/git_providers/gitlab_provider.py @@ -143,6 +143,11 @@ def __init__(self, note, mr_web_url: Optional[str] = None): self.anchor_time = max(candidates) if candidates else None self.html_url = f"{mr_web_url}#note_{self.id}" if mr_web_url else "" +# GitLab project access levels: No access = 0, Minimal Access = 5, Guest = 10, +# Reporter = 20, Developer = 30, Maintainer = 40, Owner = 50. Reading a private or +# internal project's repository source requires at least Reporter-level access. +_GITLAB_ACCESS_LEVEL_REPORTER = 20 + class GitLabProvider(GitProvider): def __init__(self, merge_request_url: Optional[str] = None, incremental: Optional[bool] = False): @@ -414,8 +419,7 @@ def _get_merge_request_changes(self) -> dict: return changes def is_supported(self, capability: str) -> bool: - if capability in ['create_inline_comment', 'publish_inline_comments', - 'publish_file_comments']: # gfm_markdown is supported in gitlab ! + if capability in ['create_inline_comment', 'publish_inline_comments']: # gfm_markdown is supported in gitlab ! return False if capability == "push_code" and get_settings().config.restricted_mode: return False @@ -424,6 +428,9 @@ def is_supported(self, capability: str) -> bool: def supports_incremental_kind(self, kind: str) -> bool: return kind in self._INCREMENTAL_ANCHOR_PREFIXES + def supports_issue_reference_tickets(self) -> bool: + return True + def _get_project_path_from_pr_or_issue_url(self, pr_or_issue_url: str) -> str: repo_project_path = None if 'issues' in pr_or_issue_url: @@ -732,7 +739,7 @@ def _get_own_user_id(self) -> Optional[int]: def get_pr_file_content(self, file_path: str, branch: str) -> str: try: - file_obj = self.gl.projects.get(self.id_project).files.get(file_path, branch) + file_obj = self.gl.projects.get(self.id_project, lazy=True).files.get(file_path, branch) content = file_obj.decode() return decode_if_bytes(content) except GitlabGetError: @@ -889,6 +896,18 @@ def get_files(self) -> list: self.git_files = [c.get('new_path') for c in raw_changes if c.get('new_path')] return self.git_files + def get_pr_file_paths(self) -> list: + """Return the complete MR file set regardless of incremental review state. + + get_files() returns only the unreviewed subset once an incremental review + is active, so per-directory settings would change between commands based on + which files the review already covered. Discovery instead walks the full MR + changes, keeping both old_path and new_path so both sides of a rename apply. + """ + raw_changes = self._get_merge_request_changes().get('changes', []) + raw_changes = self._expand_submodule_changes(raw_changes) + return [c for c in raw_changes if c.get('new_path') or c.get('old_path')] + def publish_description(self, pr_title: str, pr_body: str) -> None: try: if pr_title is not None: @@ -939,25 +958,6 @@ def is_comment_authored_by_pr_agent(self, comment) -> bool: raise RuntimeError("GitLab authenticated user cannot be verified") return str(author_id) == str(own_user_id) - def publish_persistent_comment(self, pr_comment: str, - initial_header: str, - update_header: bool = True, - name='review', - final_update_message=True, - as_thread: bool = False, - identity_marker: str | None = None, - legacy_initial_header: str | None = None): - self.publish_persistent_comment_full( - pr_comment, - initial_header, - update_header, - name, - final_update_message, - as_thread=as_thread, - identity_marker=identity_marker, - legacy_initial_header=legacy_initial_header, - ) - def publish_comment(self, mr_comment: str, is_temporary: bool = False, as_thread: bool = False): if is_temporary and not get_settings().config.publish_output_progress: get_logger().debug(f"Skipping publish_comment for temporary comment: {mr_comment}") @@ -1327,9 +1327,6 @@ def publish_code_suggestions(self, code_suggestions: list) -> bool: # note that we publish suggestions one-by-one. so, if one fails, the rest will still be published return True - def publish_file_comments(self, file_comments: list) -> bool: - pass - def search_line(self, relevant_file, relevant_line_in_file): target_file = None @@ -1407,18 +1404,20 @@ def get_languages(self): def get_pr_branch(self): return self.mr.source_branch - def get_owning_namespace(self) -> str | None: + def get_owning_namespace(self, *, resolved: bool = False) -> str | None: # The top-level group of the project's path_with_namespace works on any host - # (gitlab.com or self-hosted) with no extra round trip: numeric project IDs are - # resolved to their canonical path first so the group name is still available. + # (gitlab.com or self-hosted). Resolve numeric IDs, and all identifiers when + # checking sibling access, to the canonical path before comparing groups. if not getattr(self, "gl", None) or not getattr(self, "id_project", None): return None project_id = str(self.id_project) - if project_id.isascii() and project_id.isdigit(): + if resolved or (project_id.isascii() and project_id.isdigit()): try: project_path = self.gl.projects.get(project_id).path_with_namespace except Exception as e: get_logger().warning(f"Failed to resolve canonical GitLab project path, error: {e}") + if resolved: + raise return None if not project_path: return None @@ -1450,6 +1449,63 @@ def get_repo_settings(self): get_logger().warning(f"Failed to load local .pr_agent.toml file, error: {e}") return settings_files if settings_files else "" + def get_repo_settings_tree(self, ref: str = "") -> tuple[list[str], str]: + """Recursively list every `.pr_agent.toml` at the repository default branch. + + GitLab root config is always read from the project default branch; the + per-directory layer follows the same branch so nested configs cannot read + a branch that the root does not use. ``ref`` is accepted for interface + compatibility but ignored — a future follow-up could add CONFIG_BRANCH + support here. + """ + if not getattr(self, "gl", None) or not getattr(self, "id_project", None): + return [], "" + try: + project = self.gl.projects.get(self.id_project) + resolved_ref = project.default_branch + max_pages = get_settings().config.per_directory_settings_max_tree_pages + if isinstance(max_pages, bool) or not isinstance(max_pages, int) or max_pages < 1: + get_logger().warning("Invalid per-directory tree page limit; skipping nested settings") + return [], resolved_ref + paths = [] + for page in range(1, max_pages + 1): + tree = project.repository_tree(ref=resolved_ref, recursive=True, page=page, per_page=100) + paths.extend( + item["path"] for item in tree + if item.get("type") == "blob" + and (item.get("path") or "").split("/")[-1] == ".pr_agent.toml" + ) + if len(tree) < 100: + return paths, resolved_ref + get_logger().warning("Per-directory tree page limit reached; skipping incomplete nested settings discovery") + return [], resolved_ref + except GitlabGetError as e: + if getattr(e, "response_code", None) == 404: + get_logger().debug("No repository tree found for per-directory settings; skipping") + return [], "" + raise + + def get_repo_settings_contents(self, paths: list[str], ref: str) -> dict[str, bytes]: + """Fetch raw content of per-directory settings files at *ref*.""" + if not getattr(self, "gl", None) or not getattr(self, "id_project", None): + return {} + project = self.gl.projects.get(self.id_project) + result: dict[str, bytes] = {} + for path in paths: + try: + content = project.files.get(file_path=path, ref=ref).decode() + if isinstance(content, str): + content = content.encode("utf-8") + result[path] = content + except GitlabGetError as e: + if getattr(e, "response_code", None) == 404: + get_logger().warning( + f"Per-directory settings file '{path}' not found at ref '{ref}'; skipping" + ) + else: + raise + return result + def _get_global_settings_cache_key(self, group: str) -> str: return f"gitlab:{getattr(self, 'gitlab_url', '')}:{group}" @@ -1481,6 +1537,105 @@ def get_repo_file_content(self, file_path: str, from_default_branch: bool = Fals return "" raise + def get_sibling_repo_file_content(self, repo_id: str, file_path: str, from_default_branch: bool = False): + try: + repo_id = (repo_id or "").strip().strip("/") + file_path = (file_path or "").strip().lstrip("/") + if not repo_id or not file_path: + return "" + if not self.is_sibling_repo_allowed(repo_id): + get_logger().warning(f"Ignoring sibling repo absent from the host allowlist: {repo_id}") + return "" + project = self.gl.projects.get(repo_id) + resolved_path = project.path_with_namespace + current_namespace = self.get_owning_namespace(resolved=True) + numeric_id = repo_id.isascii() and repo_id.isdigit() + identity_matches = str(project.id) == repo_id if numeric_id else resolved_path == repo_id + if (not isinstance(resolved_path, str) or "/" not in resolved_path or not identity_matches + or not current_namespace or resolved_path.split("/")[0] != current_namespace): + get_logger().warning(f"Ignoring out-of-namespace sibling repo in repo context: {repo_id}") + return "" + if not self._requester_can_read_sibling_project(project): + get_logger().warning( + f"Ignoring sibling repo context file the review requester cannot read: {repo_id}" + ) + return "" + contents = project.files.get(file_path=file_path, ref=project.default_branch).decode() + return decode_if_bytes(contents) + except GitlabGetError as e: + # A missing optional file is expected, but transient/provider failures must reach + # repo_context so the failed result is not cached as a successful empty context. + if getattr(e, "response_code", None) == 404: + return "" + raise + + def _get_review_requester_id(self) -> Optional[int]: + # Prefer the authenticated command actor when one is known: comment commands can pass + # arbitrary arguments, so sibling context must be authorized against whoever actually + # issued the command, not the MR author. Fall back to the MR author only when no actor + # is recorded (CLI runs), and fail closed without either identity. + requester_id = getattr(self, "_command_actor", None) + if requester_id: + return requester_id + mr = getattr(self, "mr", None) + author = getattr(mr, "author", None) if mr is not None else None + if isinstance(author, dict): + return author.get("id") + if author is not None: + return getattr(author, "id", None) + return None + + def _requester_can_read_sibling_project(self, project) -> bool: + # Public projects are readable by any instance user, including external users. + visibility = getattr(project, "visibility", None) + if visibility == "public": + return True + requester_id = self._get_review_requester_id() + if not requester_id: + return False + if visibility == "internal": + # Internal projects are visible to every signed-in instance user who is not an + # external user; membership is not required. Only a positive "external" verdict + # rejects here; an unknown flag falls through to the membership check and fails + # closed, since external users cannot be granted internal access as members either. + try: + user = self.gl.users.get(requester_id) + except GitlabGetError as e: + # A definitive 404 (no such user) means denial; provider/network failures are + # not a verdict and must surface as a fetch error instead of a silent skip. + if getattr(e, "response_code", None) == 404: + return False + raise + if getattr(user, "external", None) is False: + return True + try: + member = project.members_all.get(requester_id) + except GitlabGetError as e: + # A definitive 404 (not a member) means denial; auth/provider failures are not a + # verdict and must surface as a fetch error instead of silently dropping the file. + if getattr(e, "response_code", None) == 404: + return False + raise + access_level = getattr(member, "access_level", None) + # Reading repository source code requires at least Reporter-level access; Guests, + # minimal-access members, and unknown levels are rejected. Requiring positive membership + # with that level also fails closed for external users, whom GitLab forbids from + # internal projects, and for low-access members of private projects. + return isinstance(access_level, int) and access_level >= _GITLAB_ACCESS_LEVEL_REPORTER + + def get_repo_context_ref(self, from_default_branch: bool = False) -> Optional[str]: + # The MR target branch (the branch being merged into) is the cached revision; the + # project default branch is consulted when from_default_branch is requested or no MR + # target exists, mirroring get_repo_file_content. Both are branch names, so the project + # is fetched at most once per provider rather than on every repo-context read. + if not from_default_branch: + target_branch = getattr(self.mr, "target_branch", None) + if target_branch: + return target_branch + if not hasattr(self, "_repo_context_default_branch"): + self._repo_context_default_branch = self.gl.projects.get(self.id_project).default_branch + return self._repo_context_default_branch + def get_workspace_name(self): return self.id_project.split('/')[0] diff --git a/pr_agent/git_providers/plain_diff_provider.py b/pr_agent/git_providers/plain_diff_provider.py index 939175eb15..2f55d5a57a 100644 --- a/pr_agent/git_providers/plain_diff_provider.py +++ b/pr_agent/git_providers/plain_diff_provider.py @@ -7,6 +7,7 @@ from pr_agent.algo.language_handler import build_language_file_matcher from pr_agent.algo.types import FilePatchInfo +from pr_agent.algo.utils import format_pr_code_suggestions_header, show_run_details from pr_agent.config_loader import _find_repository_root, get_settings from pr_agent.git_providers.diff_parsing import parse_unified_diff, reconstruct_base_file, to_hunk_only_patch from pr_agent.git_providers.git_provider import GitProvider @@ -136,11 +137,14 @@ def publish_description(self, pr_title: str, pr_body: str): def is_supported(self, capability: str) -> bool: if capability in ["get_issue_comments", "create_inline_comment", - "publish_inline_comments", "publish_file_comments", - "get_labels"]: + "publish_inline_comments", + "get_labels", "edit_comment", "remove_comment"]: return False return True + def supports_code_suggestions_artifact(self) -> bool: + return True + def get_languages(self): # Return {language-name: percentage}, matching the hosted providers. # sort_files_by_main_languages() keys on language NAMES (it maps each @@ -186,6 +190,18 @@ def publish_code_suggestions(self, code_suggestions: list) -> bool: # them to a (non-existent) hosting platform. if not code_suggestions: return True + return self.publish_code_suggestions_artifact(code_suggestions) + + def publish_code_suggestions_artifact( + self, code_suggestions: list, artifact_footer: str = "", + no_suggestions_message: str = "No code suggestions found for the PR.") -> bool: + if not code_suggestions: + content = f"{format_pr_code_suggestions_header()}\n\n{no_suggestions_message}{artifact_footer}" + if get_settings().get("config.output_run_details", False): + content += show_run_details(self.is_supported("gfm_markdown")) + self._write_output(content) + return True + sections = ["## Code suggestions", ""] for s in code_suggestions: relevant_file = s.get("relevant_file", "") @@ -196,7 +212,8 @@ def publish_code_suggestions(self, code_suggestions: list) -> bool: sections.append(f"### {location}") sections.append(s.get("body", "")) sections.append("") - self._write_output("\n".join(sections).rstrip() + "\n") + content = "\n".join(sections).rstrip() + artifact_footer + "\n" + self._write_output(content) return True # ---- unsupported publish operations (no-op or NotImplementedError) ---- diff --git a/pr_agent/git_providers/utils.py b/pr_agent/git_providers/utils.py index 50337df987..58d43ebaa2 100644 --- a/pr_agent/git_providers/utils.py +++ b/pr_agent/git_providers/utils.py @@ -1,5 +1,8 @@ import copy +import itertools +import math import os +import posixpath import re import tempfile import tomllib @@ -12,13 +15,25 @@ from starlette_context import context from pr_agent.config_loader import get_settings -from pr_agent.config_security import REPO_HOST_ONLY_KEYS_BY_SECTION, REPO_OVERRIDABLE_KEYS_BY_HOST_SECTION +from pr_agent.config_security import ( + PER_DIRECTORY_HOST_ONLY_KEYS_BY_SECTION, + REPO_HOST_ONLY_KEYS_BY_SECTION, + REPO_OVERRIDABLE_KEYS_BY_HOST_SECTION, + REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS, +) from pr_agent.custom_merge_loader import MAX_TOML_SIZE_IN_BYTES, validate_file_security from pr_agent.git_providers import get_git_provider_with_context from pr_agent.log import get_logger _MAX_EXTRA_CONFIG_BYTES = 1 * 1024 * 1024 # 1 MB cap for a remote .toml _FETCH_TIMEOUT_SECONDS = 10 +# Hard ceiling on the number of per-directory `.pr_agent.toml` files applied per MR, +# so a wide, multi-service diff cannot trigger a burst of config fetches. +_DEFAULT_MAX_PER_DIRECTORY_SETTINGS = 20 +# Aggregate byte budget for per-directory settings retained per MR. The per-file cap is +# MAX_TOML_SIZE_IN_BYTES (the same limit the loader applies); this trusted total keeps a +# batch of contributor-controlled files from being retained and parsed in bulk. +_PER_DIRECTORY_SETTINGS_AGGREGATE_BYTES = 5 * 1024 * 1024 # Bare Windows drive-letter paths (e.g. "C:\\shared.toml", "D:/cfg.toml"). # urlparse() would otherwise interpret the drive letter as a URL scheme. _WINDOWS_DRIVE_PATH_RE = re.compile(r"^[A-Za-z]:[\\/]") @@ -225,6 +240,7 @@ def _apply_settings_from_file(path: str, label: str): def apply_repo_settings(pr_url): os.environ["AUTO_CAST_FOR_DYNACONF"] = "false" + _restore_per_directory_settings() # Apply external/shared config FIRST, before constructing the git provider: # provider initialisers (e.g. GitLabProvider reads GITLAB.PERSONAL_ACCESS_TOKEN @@ -278,22 +294,27 @@ def apply_repo_settings(pr_url): # misattribute the error to the last source; applying per-scope keeps error reporting # (and redaction) accurate and lets valid sources still take effect. for category, settings_content in _normalize_repo_settings(repo_settings): - fd, repo_settings_file = tempfile.mkstemp(suffix='.toml') - repo_settings_files.append(repo_settings_file) - if isinstance(settings_content, str): - settings_content = settings_content.encode("utf-8") - # os.fdopen takes ownership of fd (closes it) and write() writes all bytes, - # avoiding a silently-truncated file from a partial os.write. - with os.fdopen(fd, "wb") as settings_file_handle: - settings_file_handle.write(settings_content) + repo_settings_file = _write_settings_temp(settings_content, repo_settings_files) try: _apply_repo_settings_file(repo_settings_file) except Exception as e: get_logger().warning(f"Failed to apply repo {category} settings, error: {str(e)}") config_errors.append({'error': str(e), 'settings': settings_content, 'category': category}) - if config_errors: - handle_configurations_errors(config_errors, git_provider) + # Per-directory layer (monorepo support, opt-in): merge `.pr_agent.toml` files found in + # the directories the PR touches. Applied after the root config so a nearer file + # overrides a farther one; only non-critical sections are accepted (see + # REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS). + for category, settings_content in _get_per_directory_settings(git_provider): + repo_settings_file = _write_settings_temp(settings_content, repo_settings_files) + try: + _apply_repo_settings_file(repo_settings_file, repo_settings_scope="per_directory") + except Exception as e: + get_logger().warning(f"Failed to apply per-directory settings {category}, error: {str(e)}") + config_errors.append({'error': str(e), 'settings': settings_content, 'category': category}) + + if config_errors: + handle_configurations_errors(config_errors, git_provider) except Exception as e: get_logger().exception("Failed to apply repo settings", e) finally: @@ -308,12 +329,47 @@ def apply_repo_settings(pr_url): set_claude_model() -def _apply_repo_settings_file(repo_settings_file): +def _restore_per_directory_settings(): + """Remove the previous directory overlay before applying this command's trusted settings.""" + settings = get_settings() + previous = vars(settings).pop("_per_directory_original_values", {}) + for section, original_values in previous.items(): + contents = copy.deepcopy(settings.as_dict().get(section, {})) + for key, (present, value) in original_values.items(): + for current_key in list(contents): + if current_key.lower() == key: + del contents[current_key] + if present: + contents[key] = value + settings.unset(section) + if contents: + settings.set(section, contents, merge=False) + + +def _write_settings_temp(settings_content, repo_settings_files: list) -> str: + """Write a settings payload (str or bytes) to a temp .toml file registered for cleanup. + + os.fdopen takes ownership of the fd (closes it) and write() writes all bytes, + avoiding a silently-truncated file from a partial os.write. + """ + fd, repo_settings_file = tempfile.mkstemp(suffix='.toml') + repo_settings_files.append(repo_settings_file) + if isinstance(settings_content, str): + settings_content = settings_content.encode("utf-8") + with os.fdopen(fd, "wb") as settings_file_handle: + settings_file_handle.write(settings_content) + return repo_settings_file + + +def _apply_repo_settings_file(repo_settings_file, repo_settings_scope="repo"): """Load a single repo settings file and merge its allowed keys into the global settings. Enforces the per-repo host-key restrictions and logs only section names (values may contain - secrets). Raises on load/parse failure so the caller can attribute the error to the correct - settings scope (e.g. 'global' vs 'local'). + secrets). With ``repo_settings_scope="per_directory"`` the per-directory section allowlist + (REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS) is enforced on top, so a nested config committed by + any contributor cannot touch secrets/identity or deployment-critical settings that the root + config may set. Raises on load/parse failure so the caller can attribute the error to the + correct settings scope (e.g. 'global' vs 'local'). """ # Enforce the same size cap as the loader BEFORE parsing, so an oversized file can't be fully # read/parsed in-process (OOM/CPU) by the explicit validation below. @@ -338,6 +394,87 @@ def _apply_repo_settings_file(repo_settings_file): if not isinstance(contents, dict) or not contents: get_logger().debug(f"Skipping non-table or empty section: {section}") continue + if repo_settings_scope == "per_directory": + if section.lower() not in REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS: + get_logger().warning( + f"Ignoring section [{section}] from per-directory settings: only " + f"{sorted(REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS)} may be set per directory" + ) + continue + per_dir_allowed_keys = REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS[section.lower()] + if isinstance(per_dir_allowed_keys, frozenset): + rejected = [k for k in contents if k.lower() not in per_dir_allowed_keys] + if rejected: + get_logger().warning( + f"Ignoring non-overridable key(s) {rejected} in section [{section}] from " + f"per-directory settings; only {sorted(per_dir_allowed_keys)} may be set here" + ) + contents = {k: v for k, v in contents.items() if k.lower() in per_dir_allowed_keys} + if not contents: + continue + if section.lower() == "pr_description": + normalized_contents = {key.lower(): value for key, value in contents.items()} + threshold = normalized_contents.get("collapsible_file_list_threshold") + if ( + "collapsible_file_list_threshold" in normalized_contents + and ( + isinstance(threshold, bool) + or not isinstance(threshold, int) + or not 0 <= threshold <= 1000 + ) + ): + get_logger().warning( + "Ignoring invalid collapsible_file_list_threshold from per-directory settings; expected an " + "integer between 0 and 1000" + ) + contents = { + key: value for key, value in contents.items() + if key.lower() != "collapsible_file_list_threshold" + } + if not contents: + continue + per_dir_host_only_keys = PER_DIRECTORY_HOST_ONLY_KEYS_BY_SECTION.get(section.lower(), frozenset()) + rejected = [k for k in contents if k.lower() in per_dir_host_only_keys] + if rejected: + get_logger().warning( + f"Ignoring host-only key(s) {rejected} in section [{section}] from per-directory settings" + ) + contents = {k: v for k, v in contents.items() if k.lower() not in per_dir_host_only_keys} + if not contents: + continue + if section.lower() == "config": + normalized_contents = {key.lower(): value for key, value in contents.items()} + invalid_config_keys = [ + key for key in ("model", "model_weak", "model_reasoning", "response_language") + if key in normalized_contents + and (not isinstance(normalized_contents[key], str) or not normalized_contents[key].strip()) + ] + if invalid_config_keys: + get_logger().warning( + f"Ignoring non-string or empty setting(s) {invalid_config_keys} from per-directory settings" + ) + contents = { + key: value for key, value in contents.items() if key.lower() not in invalid_config_keys + } + if not contents: + continue + invalid_temperature = ( + "temperature" in normalized_contents + and ( + isinstance(normalized_contents["temperature"], bool) + or not isinstance(normalized_contents["temperature"], (int, float)) + or not math.isfinite(normalized_contents["temperature"]) + or not 0 <= normalized_contents["temperature"] <= 2 + ) + ) + if invalid_temperature: + get_logger().warning( + "Ignoring invalid temperature setting from per-directory settings; expected a finite number " + "between 0 and 2" + ) + contents = {key: value for key, value in contents.items() if key.lower() != "temperature"} + if not contents: + continue allowed_keys = REPO_OVERRIDABLE_KEYS_BY_HOST_SECTION.get(section.lower()) if allowed_keys is not None: rejected = [k for k in contents if k.lower() not in allowed_keys] @@ -360,7 +497,20 @@ def _apply_repo_settings_file(repo_settings_file): if not contents: continue section_dict = copy.deepcopy(get_settings().as_dict().get(section.upper(), {})) + if repo_settings_scope == "per_directory": + previous = vars(get_settings()).setdefault("_per_directory_original_values", {}) + original_values = previous.setdefault(section.upper(), {}) + normalized = {key.lower(): value for key, value in section_dict.items()} + for key in contents: + original_values.setdefault( + key.lower(), (key.lower() in normalized, copy.deepcopy(normalized.get(key.lower())))) for key, value in contents.items(): + # Dynaconf looks up keys case-insensitively, so replacing the existing key + # (whatever its casing) keeps the newer value from a nearer/sibling config + # deterministic instead of leaving an "canonical-cased" duplicate beside it. + for existing_key in list(section_dict): + if existing_key.lower() == key.lower(): + del section_dict[existing_key] section_dict[key] = value get_settings().unset(section) get_settings().set(section, section_dict, merge=False) @@ -381,6 +531,261 @@ def _normalize_repo_settings(repo_settings): return repo_settings +def _get_config_branch() -> str: + """Resolve the branch repo configuration is read from. + + Mirrors the per-provider resolution (CONFIG.CONFIG_BRANCH / PR_AGENT_CONFIG_BRANCH); + returns "" when unset, meaning the provider reads from its own default branch. + """ + settings_branch = get_settings().get("CONFIG.CONFIG_BRANCH", None) + settings_branch = settings_branch.strip() if isinstance(settings_branch, str) else "" + env_branch = (os.environ.get("PR_AGENT_CONFIG_BRANCH") or "").strip() + return settings_branch or env_branch + + +def _get_changed_file_paths(git_provider) -> list[str]: + """Return the repository-relative paths the PR/MR touches. + + Includes rename metadata (old_path / previous_filename) when the provider + surfaces it, so ancestor configs for both sides of a move apply. Prefers the + complete PR file set (get_pr_file_paths) over get_files(), which many + providers narrow to the unreviewed subset while an incremental review is + active; discovery must always see the whole PR so applied settings do not + change between commands. Tolerates each provider's entry shape (str, dict + keyed by new_path/filename/path, or an object with .filename/.new_path). A + failure to list files degrades to no per-directory configs rather than + failing the request. + """ + try: + full_paths = getattr(git_provider, "get_pr_file_paths", None) + files = full_paths() if full_paths is not None else git_provider.get_files() + except Exception as e: + get_logger().warning(f"Failed to list changed files for per-directory settings: {e}") + return [] + paths = [] + for entry in files or []: + for name in _entry_path_names(entry): + if name not in paths: + paths.append(name) + return paths + + +def _entry_path_names(entry) -> list[str]: + """Extract any repository-relative path names an entry carries (new and old).""" + if isinstance(entry, str): + return [entry.strip()] if entry.strip() else [] + if isinstance(entry, dict): + keys = ("new_path", "filename", "path", "old_path", "previous_filename") + else: + keys = ("filename", "new_path", "old_path", "previous_filename") + names = [] + for key in keys: + value = entry.get(key) if isinstance(entry, dict) else getattr(entry, key, None) + if isinstance(value, str) and value.strip() and value != "/dev/null": + names.append(value.strip()) + return names + + +def _get_per_directory_settings(git_provider) -> list: + """Resolve the `.pr_agent.toml` files for the directories an MR touches. + + Walks up from each changed file toward the repository root and keeps every + ancestor config whose directory has one (the root `.pr_agent.toml` is already + applied through get_repo_settings()). Results are returned ordered + shallowest-directory-first so a nearer (more specific) file overrides a farther + one on shared keys; list/dict values replace rather than concatenate, matching + the loader. + + With a diff spanning several sibling directories (e.g. ``services/auth`` and + ``services/billing``), all of their configs apply to the whole MR: an explicit + deterministic tie-break orders equal-depth directories by path, so the + lexicographically-last sibling wins on a shared key. That is only meaningful as + a fallback, so overlapping keys between same-depth configs applied to one MR are + reported with a warning naming the conflict and the winner. + + Returns [] when the feature is disabled, the provider lacks per-directory support, + the tree has no nested configs, or nothing is crossed. + """ + settings = get_settings() + if not settings.config.get("enable_per_directory_settings", False): + return [] + config_branch = _get_config_branch() + tree_method = getattr(git_provider, "get_repo_settings_tree", None) + if tree_method is None: + return [] + try: + tree_paths, resolved_ref = tree_method(config_branch) + except Exception as e: + get_logger().warning(f"Failed to list per-directory settings candidates: {e}") + return [] + # The root `.pr_agent.toml` (directory '') is already applied through + # get_repo_settings(); only nested directories add per-directory behavior. + config_dirs = { + posixpath.dirname(p) + for p in tree_paths + if posixpath.basename(p) == ".pr_agent.toml" and posixpath.dirname(p) + } + if not config_dirs: + return [] + changed_paths = _get_changed_file_paths(git_provider) + if not changed_paths: + return [] + + try: + max_configs = int(settings.config.get( + "per_directory_settings_max_files", _DEFAULT_MAX_PER_DIRECTORY_SETTINGS)) + except (TypeError, ValueError): + max_configs = _DEFAULT_MAX_PER_DIRECTORY_SETTINGS + max_configs = max(1, max_configs) + crossed = set() + for changed_path in changed_paths: + directory = posixpath.dirname(changed_path) + while directory: + if directory in config_dirs: + crossed.add(directory) + parent = posixpath.dirname(directory) + if parent == directory: + break + directory = parent + + if len(crossed) > max_configs: + get_logger().warning( + f"{len(crossed)} per-directory .pr_agent.toml files apply to this PR; " + f"applying at most {max_configs}, preferring shallower files and keeping " + f"later-path winners within a partially included depth" + ) + # Shallowest (closest to root) first: later files override earlier ones, so the + # nearest directory wins on scalar keys. Equal-depth siblings are ordered by + # path, so the lexicographically-last directory wins deterministically (any + # overlap is surfaced by _warn_on_sibling_key_conflicts). When the cap cuts + # through an equal-depth group, the lexicographically-later (winning) entries + # are retained and the earlier ones dropped, so the documented winner still + # participates in the merge. + ordered = sorted(crossed, key=lambda directory: (directory.count("/"), directory)) + if len(ordered) > max_configs: + kept = [] + budget = max_configs + for _, siblings in itertools.groupby(ordered, key=lambda directory: directory.count("/")): + siblings = list(siblings) + if budget >= len(siblings): + kept.extend(siblings) + budget -= len(siblings) + continue + kept.extend(siblings[len(siblings) - budget:]) + break + ordered = kept + if not ordered: + return [] + + contents_method = getattr(git_provider, "get_repo_settings_contents", None) + if contents_method is None: + return [] + config_paths = [f"{directory}/.pr_agent.toml" for directory in ordered] + contents = contents_method(config_paths, resolved_ref) or {} + # Reject oversized or over-budget payloads before any conflict inspection or + # retention: files are contributor-controlled and unvalidated at fetch time, so a + # batch of them must not be held and parsed wholesale in one worker. + contents = _size_bounded_per_directory_contents(config_paths, contents) + _warn_on_sibling_key_conflicts(ordered, contents) + resolved = [] + for path in config_paths: + content = contents.get(path) + if content is None: + continue + resolved.append((path, content)) + if resolved: + get_logger().info( + f"Applying {len(resolved)} per-directory settings file(s): " + f"{sorted(entry[0] for entry in resolved)}" + ) + return resolved + + +def _size_bounded_per_directory_contents(config_paths: list[str], contents: dict[str, bytes]) -> dict[str, bytes]: + """Drop per-directory payloads that exceed the size caps before any parsing or retention. + + Per-directory files are contributor-controlled and unvalidated at fetch time. Payloads + larger than ``MAX_TOML_SIZE_IN_BYTES`` are skipped, and retention stops once the + aggregate byte budget is exhausted, so a batch of oversized nested files cannot be kept + and parsed wholesale. Files already downloaded by the provider are unavoidable, but they + are not retained, parsed, or logged beyond a bounded warning. + """ + bounded: dict[str, bytes] = {} + total_bytes = 0 + for path in config_paths: + content = contents.get(path) + if content is None: + continue + size = len(content) + if size > MAX_TOML_SIZE_IN_BYTES: + get_logger().warning( + f"Per-directory settings file '{path}' is {size} bytes " + f"(> {MAX_TOML_SIZE_IN_BYTES}); skipping it" + ) + continue + if total_bytes + size > _PER_DIRECTORY_SETTINGS_AGGREGATE_BYTES: + get_logger().warning( + f"Per-directory settings aggregate exceeds " + f"{_PER_DIRECTORY_SETTINGS_AGGREGATE_BYTES} bytes; skipping further files" + ) + break + bounded[path] = content + total_bytes += size + return bounded + + +def _warn_on_sibling_key_conflicts(ordered: list[str], contents: dict[str, bytes]) -> list[tuple[str, str, str]]: + """Log (section, key) collisions between equal-depth sibling configs. + + Equal-depth directories are applied in path order (see _get_per_directory_settings), + so the lexicographically-last sibling wins on a shared key. Surface any such overlap + instead of letting it stay an arbitrary-looking silent choice. Returns the detected + conflicts as ``(section, key, winning_directory)`` for tests; parsing failures are + ignored here, since the real file is validated when applied and then reported + through handle_configurations_errors(). + """ + conflicts: list[tuple[str, str, str]] = [] + grouped = itertools.groupby(sorted(ordered, key=lambda d: d.count("/")), + key=lambda d: d.count("/")) + for _, group in grouped: + group = list(group) + if len(group) < 2: + continue + parsed = {} + for directory in group: + content = contents.get(f"{directory}/.pr_agent.toml") + if not content: + continue + if len(content) > MAX_TOML_SIZE_IN_BYTES: + get_logger().warning( + f"Per-directory settings file '{directory}/.pr_agent.toml' exceeds " + f"{MAX_TOML_SIZE_IN_BYTES} bytes; skipping sibling-conflict inspection" + ) + continue + try: + parsed[directory] = tomllib.loads(content.decode("utf-8")) + except Exception: + continue + key_owners: dict[tuple[str, str], list[str]] = {} + for directory, data in parsed.items(): + for section, table in data.items(): + if not isinstance(table, dict): + continue + for key in table: + key_owners.setdefault((section.lower(), key.lower()), []).append(directory) + for (section, key), owners in sorted(key_owners.items()): + if len(owners) < 2: + continue + winner_directory = owners[-1] + get_logger().warning( + f"Per-directory settings at the same depth set the same key " + f"'{section}.{key}': sibling configs {sorted(owners)} all set it " + f"and '{winner_directory}/.pr_agent.toml' wins (later path overrides)" + ) + conflicts.append((section, key, winner_directory)) + return conflicts + + def handle_configurations_errors(config_errors, git_provider): try: if not any(config_errors): @@ -393,7 +798,7 @@ def handle_configurations_errors(config_errors, git_provider): header = f"❌ **PR-Agent failed to apply '{config_type}' repo settings**" body = ( f"{header}\n\nThe configuration file needs to be a valid " - "[TOML](https://qodo-merge-docs.qodo.ai/usage-guide/configuration_options/), please fix it.\n\n" + "[TOML](https://docs.pr-agent.ai/usage-guide/configuration_options/), please fix it.\n\n" ) body += f"___\n\n**Error message:**\n`{err_message}`\n\n" if config_type == "global": diff --git a/pr_agent/mosaico/executor.py b/pr_agent/mosaico/executor.py index 047ebbe3a5..42183e280e 100644 --- a/pr_agent/mosaico/executor.py +++ b/pr_agent/mosaico/executor.py @@ -26,7 +26,6 @@ from a2a.types import Part, Task, TaskState, TaskStatus from starlette_context import context as sctx -from pr_agent.algo import normalize_litellm_model from pr_agent.config_loader import get_settings, global_settings from pr_agent.log import get_logger from pr_agent.mosaico.dispatch import route_and_run_result @@ -106,37 +105,22 @@ async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None async def health_check() -> str: - """LLM-connectivity probe for /health. NO-RETRY: bypasses pr-agent's retry-wrapped - LiteLLMAIHandler.chat_completion (which would retry-hang a down LLM) and issues a - single litellm.acompletion, after applying the MOSAICO LLM settings.""" + """LLM-connectivity probe for /health. Bypasses pr-agent's retry-wrapped + LiteLLMAIHandler.chat_completion and issues a single isolated completion, + after applying the MOSAICO LLM settings.""" try: import litellm - # Construct the handler purely for its side effect of applying pr-agent's LLM - # config (api_base/key/callbacks/etc.) onto the litellm module — do NOT call its - # retry-wrapped chat_completion. + # Construct the handler to snapshot request-local provider credentials and + # routing without using its retry-wrapped chat_completion. from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler + handler = LiteLLMAIHandler() model = get_settings().get("CONFIG.MODEL", None) if not model: return "Unhealthy: no model configured" - - custom_llm_provider = str( - getattr(get_settings().litellm, "custom_llm_provider", "") or "" - ).strip().lower() - kwargs = { - "model": normalize_litellm_model(model, custom_llm_provider), - "messages": [{"role": "system", "content": "Say ping"}], - "max_tokens": 10, - "timeout": 10, - } - if custom_llm_provider: - kwargs["custom_llm_provider"] = custom_llm_provider - if getattr(handler, "api_base", None): - kwargs["api_base"] = handler.api_base - - await litellm.acompletion(**kwargs) + await handler.probe_completion(model, _completion=litellm.acompletion) return "OK" except Exception as e: get_logger().warning(f"MOSAICO health_check unhealthy: {e}") diff --git a/pr_agent/py.typed b/pr_agent/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pr_agent/servers/azuredevops_server_webhook.py b/pr_agent/servers/azuredevops_server_webhook.py index 3e0eaac012..50a1ccbe34 100644 --- a/pr_agent/servers/azuredevops_server_webhook.py +++ b/pr_agent/servers/azuredevops_server_webhook.py @@ -15,7 +15,6 @@ from starlette import status from starlette.background import BackgroundTasks from starlette.middleware import Middleware -from starlette.requests import Request from starlette.responses import JSONResponse from starlette_context import context from starlette_context.middleware import RawContextMiddleware @@ -28,6 +27,7 @@ from pr_agent.git_providers.utils import apply_repo_settings from pr_agent.log import LoggingFormat, get_logger, setup_logger from pr_agent.servers.utils import basic_auth_matches, get_pr_commands +from pr_agent.telemetry.prometheus import attach_metrics_endpoint, prometheus_metrics_enabled setup_logger(fmt=LoggingFormat.JSON, level=get_settings().get("CONFIG.LOG_LEVEL", "DEBUG")) security = HTTPBasic(auto_error=False) @@ -180,7 +180,7 @@ def handle_line_comment(body: str, thread_id: int, comment_id: int, provider: Az # currently only basic auth is supported with azure webhooks # for this reason, https must be enabled to ensure the credentials are not sent in clear text -def authorize(credentials: HTTPBasicCredentials = Depends(security)): +def authorize(credentials: HTTPBasicCredentials = Depends(security)): # noqa: B008 if not WEBHOOK_USERNAME and not WEBHOOK_PASSWORD: return if not WEBHOOK_USERNAME or not WEBHOOK_PASSWORD: @@ -302,6 +302,8 @@ async def root(): def start(): app = FastAPI(middleware=[Middleware(RawContextMiddleware)]) + if prometheus_metrics_enabled(): + attach_metrics_endpoint(router) app.include_router(router) uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "3000"))) diff --git a/pr_agent/servers/bitbucket_app.py b/pr_agent/servers/bitbucket_app.py index 8f28fea7e2..056cf79a37 100644 --- a/pr_agent/servers/bitbucket_app.py +++ b/pr_agent/servers/bitbucket_app.py @@ -196,6 +196,9 @@ async def _perform_commands_bitbucket(commands_conf: str, agent: PRAgent, api_ur if commands_conf == "pr_commands" else get_settings().get(f"bitbucket_app.{commands_conf}", {}) ) + if not commands: + get_logger().info(f"No {commands_conf} configured, skipping auto commands") + return get_settings().set("config.is_auto_command", True) if commands_conf == "push_commands": is_valid_push = await _validate_time_from_last_commit_to_pr_update(data) @@ -381,8 +384,7 @@ async def inner(): with get_logger().contextualize(**log_context): if get_identity_provider().verify_eligibility("bitbucket", sender_id, pr_url) is not Eligibility.NOT_ELIGIBLE: - if get_pr_commands("bitbucket_app"): - await _perform_commands_bitbucket("pr_commands", agent, pr_url, log_context, data) + await _perform_commands_bitbucket("pr_commands", agent, pr_url, log_context, data) elif event == "pullrequest:updated": # PR updated, might be from a push (we will validate this later) pr_url = data["data"]["pullrequest"]["links"]["html"]["href"] log_context["api_url"] = pr_url @@ -391,9 +393,7 @@ async def inner(): with get_logger().contextualize(**log_context): if get_identity_provider().verify_eligibility("bitbucket", sender_id, pr_url) is not Eligibility.NOT_ELIGIBLE: - - if get_settings().get("bitbucket_app.push_commands"): - await _perform_commands_bitbucket("push_commands", agent, pr_url, log_context, data) + await _perform_commands_bitbucket("push_commands", agent, pr_url, log_context, data) elif event == "pullrequest:comment_created": pr_url = data["data"]["pullrequest"]["links"]["html"]["href"] log_context["api_url"] = pr_url @@ -409,7 +409,7 @@ async def inner(): return "OK" @router.get("/webhook") -async def handle_github_webhooks(request: Request, response: Response): +async def handle_webhook_health(request: Request, response: Response): return "Webhook server online!" @router.post("/installed") diff --git a/pr_agent/servers/gerrit_server.py b/pr_agent/servers/gerrit_server.py index b136570eb5..9dbcfe7dad 100644 --- a/pr_agent/servers/gerrit_server.py +++ b/pr_agent/servers/gerrit_server.py @@ -21,7 +21,7 @@ security = HTTPBasic(auto_error=False) -def authorize(credentials: HTTPBasicCredentials = Depends(security)): +def authorize(credentials: HTTPBasicCredentials = Depends(security)): # noqa: B008 """Require the configured webhook credentials on every call.""" gerrit_settings = get_settings().get("gerrit", {}) username = gerrit_settings.get("webhook_username", None) if gerrit_settings else None diff --git a/pr_agent/servers/gitea_app.py b/pr_agent/servers/gitea_app.py index 4fb187f582..aa134c7b73 100644 --- a/pr_agent/servers/gitea_app.py +++ b/pr_agent/servers/gitea_app.py @@ -14,6 +14,7 @@ from pr_agent.git_providers.utils import apply_repo_settings from pr_agent.log import LoggingFormat, get_logger, setup_logger from pr_agent.servers.utils import get_pr_commands, push_trigger_slot, verify_signature +from pr_agent.telemetry.prometheus import attach_metrics_endpoint, prometheus_metrics_enabled # Setup logging and router setup_logger(fmt=LoggingFormat.JSON, level=get_settings().get("CONFIG.LOG_LEVEL", "DEBUG")) @@ -92,6 +93,10 @@ async def handle_pr_event(body: Dict[str, Any], event: str, action: str, agent: if not api_url: return + apply_repo_settings(api_url) + if not should_process_pr_logic(body): + return {} + # Handle PR based on action if action in ["opened", "reopened"]: # commands = get_settings().get("gitea.pr_commands", []) @@ -129,12 +134,9 @@ async def handle_comment_event(body: Dict[str, Any], event: str, action: str, ag await agent.handle_request(pr_url, comment_body) async def _perform_commands_gitea(commands_conf: str, agent: PRAgent, body: dict, api_url: str): - apply_repo_settings(api_url) if commands_conf == "pr_commands" and get_settings().config.disable_auto_feedback: # auto commands for PR, and auto feedback is disabled get_logger().info(f"Auto feedback is disabled, skipping auto commands for PR {api_url=}") return - if not should_process_pr_logic(body): # Here we already updated the configuration with the repo settings - return {} commands = ( get_pr_commands("gitea") if commands_conf == "pr_commands" @@ -212,6 +214,8 @@ def should_process_pr_logic(body) -> bool: # FastAPI app setup middleware = [Middleware(RawContextMiddleware)] +if prometheus_metrics_enabled(): + attach_metrics_endpoint(router) app = FastAPI(middleware=middleware) app.include_router(router) diff --git a/pr_agent/servers/github_app.py b/pr_agent/servers/github_app.py index 16a3f7de96..06478f3a7c 100644 --- a/pr_agent/servers/github_app.py +++ b/pr_agent/servers/github_app.py @@ -1,6 +1,7 @@ import copy import os import re +import time import uuid from typing import Any, Dict, Tuple @@ -18,7 +19,8 @@ from pr_agent.identity_providers import get_identity_provider from pr_agent.identity_providers.identity_provider import Eligibility from pr_agent.log import LoggingFormat, get_logger, setup_logger -from pr_agent.servers.utils import get_pr_commands, push_trigger_slot, verify_signature +from pr_agent.servers.utils import DefaultDictWithTimeout, get_pr_commands, push_trigger_slot, verify_signature +from pr_agent.telemetry.prometheus import attach_metrics_endpoint, prometheus_metrics_enabled setup_logger(fmt=LoggingFormat.JSON, level=get_settings().get("CONFIG.LOG_LEVEL", "DEBUG")) base_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) @@ -29,6 +31,12 @@ else: build_number = "unknown" router = APIRouter() +_WEBHOOK_DELIVERY_TTL = get_settings().github_app.push_trigger_pending_tasks_ttl +_completed_webhook_deliveries = DefaultDictWithTimeout( + float, + ttl=_WEBHOOK_DELIVERY_TTL, + update_key_time_on_get=False, +) @router.post("/api/v1/github_webhooks") @@ -46,7 +54,12 @@ async def handle_github_webhooks(background_tasks: BackgroundTasks, request: Req context["installation_id"] = installation_id context["settings"] = copy.deepcopy(global_settings) context["git_provider"] = {} - background_tasks.add_task(handle_request, body, event=request.headers.get("X-GitHub-Event", None)) + background_tasks.add_task( + handle_request, + body, + event=request.headers.get("X-GitHub-Event", None), + delivery_id=request.headers.get("X-GitHub-Delivery", None), + ) return {} @@ -127,6 +140,8 @@ async def handle_comments_on_pr(body: Dict[str, Any], log_context["api_url"] = api_url comment_id = body.get("comment", {}).get("id") provider = get_git_provider_with_context(pr_url=api_url) + if isinstance(sender, str) and sender and hasattr(provider, "set_command_actor"): + provider.set_command_actor(sender) with get_logger().contextualize(**log_context): if get_identity_provider().verify_eligibility("github", sender_id, api_url) is not Eligibility.NOT_ELIGIBLE: get_logger().info(f"Processing comment on PR {api_url=}, comment_body={comment_body}") @@ -136,6 +151,7 @@ async def handle_comments_on_pr(body: Dict[str, Any], # Optional, and disabled by default: tell the author how the command ended without # adding another comment to the thread. provider.react_to_outcome(comment_id, bool(succeeded)) + return succeeded else: get_logger().info(f"User {sender=} is not eligible to process comment on PR {api_url=}") @@ -156,7 +172,7 @@ async def handle_new_pr_opened(body: Dict[str, Any], # logic to ignore PRs with specific titles (e.g. "[Auto] ...") apply_repo_settings(api_url) if get_identity_provider().verify_eligibility("github", sender_id, api_url) is not Eligibility.NOT_ELIGIBLE: - await _perform_auto_commands_github("pr_commands", agent, body, api_url, log_context) + return await _perform_auto_commands_github("pr_commands", agent, body, api_url, log_context) else: get_logger().info(f"User {sender=} is not eligible to process PR {api_url=}") @@ -228,7 +244,7 @@ async def handle_pull_request_review_submitted(body: Dict[str, Any], return {} if get_identity_provider().verify_eligibility("github", sender_id, api_url) is not Eligibility.NOT_ELIGIBLE: - await _perform_auto_commands_github("review_commands", agent, body, api_url, log_context) + return await _perform_auto_commands_github("review_commands", agent, body, api_url, log_context) else: get_logger().info(f"User {sender=} is not eligible to process review on PR {api_url=}") @@ -266,7 +282,7 @@ async def handle_push_trigger_for_new_commits(body: Dict[str, Any], return {} if get_identity_provider().verify_eligibility("github", sender_id, api_url) is not Eligibility.NOT_ELIGIBLE: get_logger().info(f"Performing incremental review for {api_url=} because of {event=} and {action=}") - await _perform_auto_commands_github("push_commands", agent, body, api_url, log_context) + return await _perform_auto_commands_github("push_commands", agent, body, api_url, log_context) def handle_closed_pr(body, event, action, log_context): @@ -375,19 +391,8 @@ def should_process_pr_logic(body) -> bool: return True -async def handle_request(body: Dict[str, Any], event: str): - """ - Handle incoming GitHub webhook requests. - - Args: - body: The request body. - event: The GitHub event type (e.g. "pull_request", "issue_comment", etc.). - """ - action = body.get("action") # "created", "opened", "reopened", "ready_for_review", "review_requested", "synchronize" - get_logger().debug(f"Handling request with event: {event}, action: {action}") - if not action: - get_logger().debug("No action found in request body, exiting handle_request") - return {} +async def _dispatch_request(body: Dict[str, Any], event: str, action: str): + """Return False when a command fails, preserving ignored events as completed work.""" agent = PRAgent() log_context, sender, sender_id, sender_type = get_log_context(body, event, action, build_number) @@ -406,22 +411,22 @@ async def handle_request(body: Dict[str, Any], event: str): # handle submitted pull request reviews elif event == 'pull_request_review' and action == 'submitted': get_logger().debug('Request body', artifact=body, event=event) - await handle_pull_request_review_submitted( + return await handle_pull_request_review_submitted( body, event, sender, sender_id, sender_type, action, log_context, agent ) # handle comments on PRs elif action == 'created': get_logger().debug('Request body', artifact=body, event=event) - await handle_comments_on_pr(body, event, sender, sender_id, action, log_context, agent) + return await handle_comments_on_pr(body, event, sender, sender_id, action, log_context, agent) # handle new PRs elif event == 'pull_request' and action != 'synchronize' and action != 'closed': get_logger().debug('Request body', artifact=body, event=event) - await handle_new_pr_opened(body, event, sender, sender_id, action, log_context, agent) + return await handle_new_pr_opened(body, event, sender, sender_id, action, log_context, agent) elif event == "issue_comment" and 'edited' in action: pass # handle_checkbox_clicked # handle pull_request event with synchronize action - "push trigger" for new commits elif event == 'pull_request' and action == 'synchronize': - await handle_push_trigger_for_new_commits(body, event, sender,sender_id, action, log_context, agent) + return await handle_push_trigger_for_new_commits(body, event, sender,sender_id, action, log_context, agent) elif event == 'pull_request' and action == 'closed': if get_settings().get("CONFIG.ANALYTICS_FOLDER", ""): handle_closed_pr(body, event, action, log_context) @@ -430,6 +435,38 @@ async def handle_request(body: Dict[str, Any], event: str): return {} +async def handle_request(body: Dict[str, Any], event: str, delivery_id: str | None = None): + """ + Handle incoming GitHub webhook requests. + + Args: + body: The request body. + event: The GitHub event type (e.g. "pull_request", "issue_comment", etc.). + delivery_id: GitHub's stable identifier for this webhook delivery and its redeliveries. + """ + action = body.get("action") # "created", "opened", "reopened", "ready_for_review", "review_requested", "synchronize" + get_logger().debug(f"Handling request with event: {event}, action: {action}") + if not action: + get_logger().debug("No action found in request body, exiting handle_request") + return {} + settings = get_settings() + if not delivery_id or not settings.get("GITHUB_APP.WEBHOOK_DELIVERY_DEDUPLICATION", False): + await _dispatch_request(body, event, action) + return {} + + async with push_trigger_slot( + f"github-webhook-delivery:{delivery_id}", + allow_backlog=False, + ttl=_WEBHOOK_DELIVERY_TTL, + ) as proceed: + if not proceed or _completed_webhook_deliveries[delivery_id] > time.monotonic(): + return {} + succeeded = await _dispatch_request(body, event, action) + if succeeded is not False: + _completed_webhook_deliveries[delivery_id] = time.monotonic() + _WEBHOOK_DELIVERY_TTL + return {} + + def handle_line_comments(body: Dict, comment_body: [str, Any]): if not comment_body: return "" @@ -505,13 +542,17 @@ async def _perform_auto_commands_github(commands_conf: str, agent: PRAgent, body get_logger().info(f"No {commands_conf} configured, skipping auto commands") return get_settings().set("config.is_auto_command", True) + succeeded = True for command in commands: try: new_command = prepare_command(command) get_logger().info(f"{commands_conf}. Performing auto command '{new_command}', for {api_url=}") - await agent.handle_request(api_url, new_command) + if await agent.handle_request(api_url, new_command) is False: + succeeded = False except Exception as e: get_logger().error(f"Failed to perform command {command}: {e}") + succeeded = False + return succeeded @router.get("/") @@ -524,6 +565,8 @@ async def root(): get_settings().set("GITHUB.DEPLOYMENT_TYPE", "app") # get_settings().set("CONFIG.PUBLISH_OUTPUT_PROGRESS", False) middleware = [Middleware(RawContextMiddleware)] +if prometheus_metrics_enabled(): + attach_metrics_endpoint(router) app = FastAPI(middleware=middleware) app.include_router(router) diff --git a/pr_agent/servers/github_polling.py b/pr_agent/servers/github_polling.py index 92ad1c94bf..d65b3d5c65 100644 --- a/pr_agent/servers/github_polling.py +++ b/pr_agent/servers/github_polling.py @@ -21,6 +21,11 @@ NOTIFICATION_URL = "https://api.github.com/notifications" DEFAULT_POLLING_REQUEST_TIMEOUT = 10 MAX_POLLING_REQUEST_TIMEOUT = 60 +POLLING_CAPACITY_CHECK_INTERVAL = 0.25 + + +class _PollingWorkerStartError(RuntimeError): + """Stop dispatch when child startup leaves process state uncertain.""" def _get_polling_request_timeout() -> float: @@ -196,23 +201,50 @@ async def is_valid_notification(notification, headers, handled_ids, session, use return False, handled_ids -def _start_queued_processes(task_queue, max_allowed_parallel_tasks): - """Start at most max_allowed_parallel_tasks queued jobs and clear the queue. +def _reap_finished_processes(active_processes): + """Release completed workers without waiting for live ones.""" + for process in active_processes[:]: + if not process.is_alive(): + process.join(timeout=0) + process.close() + active_processes.remove(process) - Do not join the started processes; let the polling loop move on to the next - iteration without waiting for them to complete. - """ - processes = [] - for i, (func, args) in enumerate(task_queue): - if i >= max_allowed_parallel_tasks: - get_logger().error( - f"Dropping {len(task_queue) - max_allowed_parallel_tasks} tasks from polling session") - break - process = multiprocessing.Process(target=func, args=args) - processes.append(process) - process.start() - task_queue.clear() - return processes + +async def _start_queued_processes(task_queue, max_allowed_parallel_tasks, active_processes): + """Keep the batch limit, waiting for capacity shared across polling iterations.""" + if max_allowed_parallel_tasks <= 0: + raise ValueError("The polling process limit must be positive") + overflow = len(task_queue) - max_allowed_parallel_tasks + if overflow > 0: + get_logger().error(f"Dropping {overflow} tasks from polling session") + for _ in range(overflow): + task_queue.pop() + + try: + while task_queue: + _reap_finished_processes(active_processes) + if len(active_processes) >= max_allowed_parallel_tasks: + await asyncio.sleep(POLLING_CAPACITY_CHECK_INTERVAL) + continue + func, args = task_queue[0] + process = multiprocessing.Process(target=func, args=args) + try: + process.start() + except BaseException as exc: + # Treat a PID-bearing child as potentially dispatched; never retry its task. + if process.pid is not None: + active_processes.append(process) + task_queue.popleft() + else: + process.close() + if isinstance(exc, Exception): + raise _PollingWorkerStartError("Polling worker startup failed; stopping dispatch") from exc + raise + active_processes.append(process) + task_queue.popleft() + finally: + if task_queue: + get_logger().error(f"Polling dispatch stopped with {len(task_queue)} tasks not dispatched") async def polling_loop(): @@ -239,8 +271,11 @@ async def polling_loop(): if not token: raise ValueError("User token must be set to get notifications") + active_processes = [] async with aiohttp.ClientSession() as session: while True: + task_queue = deque() + dispatch_started = False try: await asyncio.sleep(5) headers = { @@ -264,7 +299,6 @@ async def polling_loop(): if not notifications: continue get_logger().info(f"Received {len(notifications)} notifications") - task_queue = deque() for notification in notifications: if not notification: continue @@ -288,14 +322,21 @@ async def polling_loop(): max_allowed_parallel_tasks = 10 if task_queue: - _start_queued_processes(task_queue, max_allowed_parallel_tasks) + dispatch_started = True + await _start_queued_processes(task_queue, max_allowed_parallel_tasks, active_processes) elif response.status != 304: print(f"Failed to fetch notifications. Status code: {response.status}") + except _PollingWorkerStartError: + raise except Exception as e: get_logger().error(f"Polling exception during processing of a notification: {e}", artifact={"traceback": traceback.format_exc()}) + finally: + if task_queue and not dispatch_started: + get_logger().error(f"Polling dispatch stopped with {len(task_queue)} tasks not dispatched") + _reap_finished_processes(active_processes) if __name__ == '__main__': diff --git a/pr_agent/servers/gitlab_webhook.py b/pr_agent/servers/gitlab_webhook.py index 71cda33e30..fffc7712c8 100644 --- a/pr_agent/servers/gitlab_webhook.py +++ b/pr_agent/servers/gitlab_webhook.py @@ -23,6 +23,7 @@ from pr_agent.log import LoggingFormat, get_logger, setup_logger from pr_agent.secret_providers import get_secret_provider, validate_secret_provider_setting from pr_agent.servers.utils import get_pr_commands, push_trigger_slot +from pr_agent.telemetry.prometheus import attach_metrics_endpoint, prometheus_metrics_enabled setup_logger(fmt=LoggingFormat.JSON, level=get_settings().get("CONFIG.LOG_LEVEL", "DEBUG")) router = APIRouter() @@ -56,6 +57,14 @@ async def handle_request(api_url: str, body: str, log_context: dict, sender_id: log_context["api_url"] = api_url log_context["app_name"] = get_settings().get("CONFIG.APP_NAME", "Unknown") + # Comment commands can pass arbitrary arguments, so sibling-repo context is authorized + # against the commenter (the command actor) instead of the MR author. Fail closed when no + # trustworthy account can be recorded (e.g. the "unknown" fallback for missing sender data). + if isinstance(sender_id, int) and sender_id: + provider = get_git_provider_with_context(pr_url=api_url) + if hasattr(provider, "set_command_actor"): + provider.set_command_actor(sender_id) + with get_logger().contextualize(**log_context): await PRAgent().handle_request(api_url, body, notify) @@ -497,6 +506,8 @@ async def root(): raise ValueError("GITLAB.URL is not set") get_settings().config.git_provider = "gitlab" middleware = [Middleware(RawContextMiddleware)] +if prometheus_metrics_enabled(): + attach_metrics_endpoint(router) app = FastAPI(middleware=middleware) app.include_router(router) diff --git a/pr_agent/servers/gunicorn_config.py b/pr_agent/servers/gunicorn_config.py index 8597c635da..37804b4229 100644 --- a/pr_agent/servers/gunicorn_config.py +++ b/pr_agent/servers/gunicorn_config.py @@ -1,8 +1,6 @@ import gc import os -# from prometheus_client import multiprocess - # Sample Gunicorn configuration file. # @@ -280,6 +278,37 @@ def when_ready(server): # object it visits. Freezing moves everything allocated so far into a permanent # generation the collector never traverses, keeping those pages shared. gc.freeze() + _prepare_prometheus() + + +def _prepare_prometheus(): + """Provision the multiprocess state dir before any worker imports prometheus_client. + + prometheus_client decides at import time whether metrics are multiprocess-capable, + so the dir must exist (and PROMETHEUS_MULTIPROC_DIR be set) in the master before the + first worker fork — workers import the client lazily and inherit the env var. + """ + from pr_agent.config_loader import get_settings + from pr_agent.telemetry.prometheus_multiproc import ensure_prometheus_multiproc_dir + from pr_agent.telemetry.types import ExporterType + + settings = get_settings() + if settings.get("OTEL.IS_ENABLED", False) and settings.get("OTEL.EXPORTER_TYPE") == ExporterType.PROMETHEUS: + ensure_prometheus_multiproc_dir(str(settings.get("OTEL.PROMETHEUS_MULTIPROC_DIR", "/tmp/pr-agent-prometheus"))) + + +def child_exit(server, worker): + """Called in the master when a worker exits; drop the worker's stale state files.""" + from pr_agent.telemetry.prometheus_multiproc import prometheus_multiproc_dir + + if not prometheus_multiproc_dir(): + return + # Imported lazily: the master only ever sees prometheus_client at exit time, never + # while importing the app under preload_app (which would break multiprocess mode in + # the forked workers). + from prometheus_client import multiprocess + + multiprocess.mark_process_dead(worker.pid) def post_fork(server, worker): diff --git a/pr_agent/servers/help.py b/pr_agent/servers/help.py index aa44601f78..058216f33c 100644 --- a/pr_agent/servers/help.py +++ b/pr_agent/servers/help.py @@ -11,7 +11,7 @@ def get_general_commands_text(): "> - **/update_changelog**: Update the changelog based on the PR's contents. \n" \ "> - **/add_docs**: Generate docstring for new components introduced in the PR. \n" \ "> - **/generate_labels**: Generate labels for the PR based on the PR's contents. \n\n" \ - ">See the [tools guide](https://pr-agent-docs.codium.ai/tools/) for more details.\n" \ + ">See the [tools guide](https://docs.pr-agent.ai/tools/) for more details.\n" \ ">To list the possible configuration parameters, add a **/config** comment. \n" return commands_text @@ -25,14 +25,14 @@ def get_general_bot_help_text(): def get_review_usage_guide(): output ="**Overview:**\n" output += (f"{COMMAND_DESCRIPTIONS['review']} " - "More feedback can be [added](https://pr-agent-docs.codium.ai/tools/review/#general-configurations) by configuring the tool.\n\n" - "The tool can be triggered [automatically](https://pr-agent-docs.codium.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) every time a new PR is opened, or can be invoked manually by commenting on any PR.\n") + "More feedback can be [added](https://docs.pr-agent.ai/tools/review/#configuration-options) by configuring the tool.\n\n" + "The tool can be triggered [automatically](https://docs.pr-agent.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) every time a new PR is opened, or can be invoked manually by commenting on any PR.\n") output +="""\ - When commenting, to edit [configurations](https://github.com/Codium-ai/pr-agent/blob/main/pr_agent/settings/configuration.toml#L23) related to the review tool (`pr_reviewer` section), use the following template: ``` /review --pr_reviewer.some_config1=... --pr_reviewer.some_config2=... ``` -- With a [configuration file](https://pr-agent-docs.codium.ai/usage-guide/configuration_options/), use the following template: +- With a [configuration file](https://docs.pr-agent.ai/usage-guide/configuration_options/), use the following template: ``` [pr_reviewer] some_config1=... @@ -40,7 +40,7 @@ def get_review_usage_guide(): ``` """ - output += "\n\nSee the review [usage page](https://pr-agent-docs.codium.ai/tools/review/) for a comprehensive guide on using this tool.\n\n" + output += "\n\nSee the review [usage page](https://docs.pr-agent.ai/tools/review/) for a comprehensive guide on using this tool.\n\n" return output @@ -50,14 +50,14 @@ def get_review_usage_guide(): def get_describe_usage_guide(): output = "**Overview:**\n" output += f"{COMMAND_DESCRIPTIONS['describe']} " - output += "The tool can be triggered [automatically](https://pr-agent-docs.codium.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) every time a new PR is opened, or can be invoked manually by commenting on a PR.\n" + output += "The tool can be triggered [automatically](https://docs.pr-agent.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) every time a new PR is opened, or can be invoked manually by commenting on a PR.\n" output += """\ When commenting, to edit [configurations](https://github.com/Codium-ai/pr-agent/blob/main/pr_agent/settings/configuration.toml#L46) related to the describe tool (`pr_description` section), use the following template: ``` /describe --pr_description.some_config1=... --pr_description.some_config2=... ``` -With a [configuration file](https://pr-agent-docs.codium.ai/usage-guide/configuration_options/), use the following template: +With a [configuration file](https://docs.pr-agent.ai/usage-guide/configuration_options/), use the following template: ``` [pr_description] some_config1=... @@ -69,7 +69,7 @@ def get_describe_usage_guide(): # automation output += "
      Enabling\\disabling automation
      \n\n" output += """\ -- When you first install the app, the [default mode](https://pr-agent-docs.codium.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) for the describe tool is: +- When you first install the app, the [default mode](https://docs.pr-agent.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) for the describe tool is: ``` pr_commands = ["/describe", ...] ``` @@ -95,7 +95,7 @@ def get_describe_usage_guide(): output += """\ The default labels of the `describe` tool are quite generic: [`Bug fix`, `Tests`, `Enhancement`, `Documentation`, `Other`]. -If you specify [custom labels](https://pr-agent-docs.codium.ai/tools/describe/#handle-custom-labels-from-the-repos-labels-page) in the repo's labels page or via configuration file, you can get tailored labels for your use cases. +If you specify [custom labels](https://docs.pr-agent.ai/tools/describe/#handle-custom-labels-from-the-repos-labels-page) in the repo's labels page or via configuration file, you can get tailored labels for your use cases. Examples for custom labels: - `Main topic:performance` - pr_agent:The main topic of this PR is performance - `New endpoint` - pr_agent:A new endpoint was added in this PR @@ -137,7 +137,7 @@ def get_describe_usage_guide(): output += "" - output += "\n\nSee the [describe usage](https://pr-agent-docs.codium.ai/tools/describe/) page for a comprehensive guide on using this tool.\n\n" + output += "\n\nSee the [describe usage](https://docs.pr-agent.ai/tools/describe/) page for a comprehensive guide on using this tool.\n\n" return output @@ -163,7 +163,7 @@ def get_ask_usage_guide(): # # output += "" - output += "\n\nSee the [ask usage](https://pr-agent-docs.codium.ai/tools/ask/) page for a comprehensive guide on using this tool.\n\n" + output += "\n\nSee the [ask usage](https://docs.pr-agent.ai/tools/ask/) page for a comprehensive guide on using this tool.\n\n" return output @@ -172,7 +172,7 @@ def get_ask_usage_guide(): def get_improve_usage_guide(): output = "**Overview:**\n" output += f"{COMMAND_DESCRIPTIONS['improve']} " - output += "The tool can be triggered [automatically](https://pr-agent-docs.codium.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) every time a new PR is opened, or can be invoked manually by commenting on a PR.\n" + output += "The tool can be triggered [automatically](https://docs.pr-agent.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) every time a new PR is opened, or can be invoked manually by commenting on a PR.\n" output += """\ - When commenting, to edit [configurations](https://github.com/Codium-ai/pr-agent/blob/main/pr_agent/settings/configuration.toml#L78) related to the improve tool (`pr_code_suggestions` section), use the following template: @@ -180,7 +180,7 @@ def get_improve_usage_guide(): /improve --pr_code_suggestions.some_config1=... --pr_code_suggestions.some_config2=... ``` -- With a [configuration file](https://pr-agent-docs.codium.ai/usage-guide/configuration_options/), use the following template: +- With a [configuration file](https://docs.pr-agent.ai/usage-guide/configuration_options/), use the following template: ``` [pr_code_suggestions] @@ -190,7 +190,7 @@ def get_improve_usage_guide(): """ - output += "\n\nSee the improve [usage page](https://pr-agent-docs.codium.ai/tools/improve/) for a comprehensive guide on using this tool.\n\n" + output += "\n\nSee the improve [usage page](https://docs.pr-agent.ai/tools/improve/) for a comprehensive guide on using this tool.\n\n" return output @@ -205,5 +205,5 @@ def get_help_docs_usage_guide(): /help_docs "..." ``` """ - output += "\n\nSee the [help_docs usage](https://pr-agent-docs.codium.ai/tools/help_docs/) page for a comprehensive guide on using this tool.\n\n" + output += "\n\nSee the [help_docs usage](https://docs.pr-agent.ai/tools/help_docs/) page for a comprehensive guide on using this tool.\n\n" return output diff --git a/pr_agent/settings/.secrets_template.toml b/pr_agent/settings/.secrets_template.toml index 065fdcb909..02e03c697c 100644 --- a/pr_agent/settings/.secrets_template.toml +++ b/pr_agent/settings/.secrets_template.toml @@ -2,7 +2,7 @@ # Copy this file to .secrets.toml in the same folder. # The minimum workable settings - set openai.key to your API key. # Set github.deployment_type to "user" and github.user_token to your GitHub personal access token. -# This will allow you to run the CLI scripts in the scripts/ folder and the github_polling server. +# This will allow you to run the pr-agent CLI and the github_polling server. # # See README for details about GitHub App deployment. @@ -18,12 +18,13 @@ key = "" # Acquire through https://platform.openai.com # OpenAI Flex Processing (optional, for cost savings) # [litellm] -# extra_body='{"processing_mode": "flex"}' +# extra_body='{"service_tier": "flex"}' # model_id = "" # Optional: Custom inference profile ID for Amazon Bedrock [pinecone] api_key = "..." -environment = "gcp-starter" +cloud = "aws" +region = "us-east-1" [qdrant] # For Qdrant Cloud or self-hosted Qdrant @@ -123,6 +124,12 @@ api_token = "" org = "" pat = "" +[jira] +# For fetching ticket context from Jira Cloud (only Jira Cloud is supported) +jira_site = "" # the "" in https://.atlassian.net (e.g. "your-org") +jira_api_email = "" # Atlassian account email +jira_api_token = "" # Atlassian API token + [azure_devops_server] # For Azure devops Server basic auth - configured in the webhook creation # Optional, uncomment if you want to use Azure devops webhooks. Value assigned when you create the webhook diff --git a/pr_agent/settings/configuration.toml b/pr_agent/settings/configuration.toml index a1ac964f58..62feb66dd0 100644 --- a/pr_agent/settings/configuration.toml +++ b/pr_agent/settings/configuration.toml @@ -21,6 +21,9 @@ log_level="DEBUG" # Configurations use_repo_settings_file=true use_global_settings_file=true +enable_per_directory_settings=false # when true, merge per-directory .pr_agent.toml files found by walking up from the PR's changed files (monorepo support). Adds bounded recursive tree discovery per MR; per-directory files may only override non-critical sections (see REPO_PER_DIRECTORY_OVERRIDABLE_SECTIONS). Nearest (deepest) directory wins on shared keys; equal-depth siblings resolve to the lexicographically-last path; when more files match than the per_directory_settings_max_files cap, shallower files are applied first and a partially capped depth keeps its later-path (winning) siblings; any overlap is logged as a warning. +per_directory_settings_max_files=20 # hard ceiling on the number of per-directory .pr_agent.toml files applied per MR (deeper configs beyond the cap are skipped with a warning) +per_directory_settings_max_tree_pages=10 # maximum GitLab recursive-tree pages (100 entries each); skip nested settings if discovery is incomplete. Root/host-controlled, independent of the settings-file cap. extra_config_url="" # optional URL or path to an additional .pr_agent.toml merged before the repo-local config; also settable via --extra_config_url or PR_AGENT_EXTRA_CONFIG_URL. See docs/docs/usage-guide/configuration_options.md#external-configuration-url. disable_auto_feedback = false enable_auto_approval = false # when true, /review may auto-approve a PR via auto_approve_logic(); that caller is currently commented out @@ -30,9 +33,14 @@ retry_same_model_on_timeout = true # when false, a timed-out call is not retried skip_keys = [] custom_reasoning_model = false # when true, disables system messages and temperature controls for models that don't support chat-style inputs response_language="en-US" # Language locales code for PR responses in ISO 3166 and ISO 639 format (e.g., "en-US", "it-IT", "zh-CN", ...) -repo_context_files = ["AGENTS.md"] # Repository-relative files (e.g. AGENTS.md, CLAUDE.md) to include as AI prompt context; set to [] to disable +repo_context_files = ["AGENTS.md"] # Repository-relative files (e.g. AGENTS.md, CLAUDE.md) to include as AI prompt context; set to [] to disable local context. A structured entry {"repo_id" = ..., "file_path" = ...} selects a sibling default-branch file from the same namespace/owner; repo_id must be in the host-issued repo_context_sibling_repos allowlist below. Reads use the sibling default branch and share repo_context_max_lines; repository settings may select entries, comment arguments cannot override this key +# The repo-context cache is keyed on the revision read: the PR target commit for target-branch +# reads, and the resolved default branch head on GitHub, so a default-branch push invalidates +# cached content within the TTL. Providers without commit resolution key on a branch name or None. repo_context_from_default_branch = true # Read repo context files from the repository default branch (trusts only default-branch content). Set to false to read from the PR target branch instead. repo_context_max_lines = 500 # Maximum total rendered lines for repo context, including wrapper tags +repo_context_sibling_repos = [] # Host-only list of approved sibling repository identifiers (GitHub owner/repo, GitLab group/project or numeric ID strings) that repo_context_files sibling entries may select. Empty disables sibling reads. Approve only repositories whose content may be disclosed in consuming PRs, because the actor check bounds who triggers a read, not who chose the target or where the output lands. Repository settings and comment arguments cannot change this list. Canonical identities and owning namespaces are checked after resolution +repo_context_max_sibling_files = 5 # Maximum number of sibling-repository files fetched per repo-context build. The fetch count is bounded separately from repo_context_max_lines so selected sibling files cannot trigger an unbounded number of cross-repository calls; sibling files still compete for the repo_context_max_lines budget. Host-only (cannot be raised by a repository's .pr_agent.toml or a comment command) and clamped to a hard ceiling of 20 fetches per build. # token limits max_description_tokens = 500 max_commits_tokens = 500 @@ -146,7 +154,7 @@ require_all_thresholds_for_incremental_review=false minimal_commits_for_incremental_review=0 minimal_minutes_for_incremental_review=0 enable_intro_text=true -enable_help_text=false # Determines whether to include help text in the PR review. Enabled by default. +enable_help_text=false # Determines whether to include help text in the PR review. enable_review_coverage_footer=true # large-diff chunking (opt-in). When the token budget leaves files out of the review, split the diff # into chunks, review each chunk, and merge the per-chunk results into one review. @@ -157,7 +165,6 @@ max_number_of_calls=3 # maximum number of chunk review calls, used only when ena publish_labels=false add_original_user_description=true generate_ai_title=false -use_bullet_points=true extra_instructions = "" enable_pr_type=true enable_pr_description=true # adds a section with an AI-generated summary of the PR changes @@ -175,7 +182,6 @@ enable_semantic_files_types=true collapsible_file_list='adaptive' # true, false, 'adaptive' collapsible_file_list_threshold=6 file_table_collapsible_open_by_default=false -inline_file_summary=false # false, true, 'table' # markers use_description_markers=false enable_large_pr_handling=true @@ -214,14 +220,12 @@ new_score_mechanism=true new_score_mechanism_th_high=9 new_score_mechanism_th_medium=7 # params for '/improve --extended' mode -auto_extended_mode=true num_code_suggestions_per_chunk=3 # Maximum suggestions retained per file after all chunks are merged; 0 disables the cap. max_suggestions_per_file = 0 max_number_of_calls = 3 parallel_calls = true -final_clip_factor = 0.8 decouple_hunks = false # self-review checkbox demand_code_suggestions_self_review=false # add a checkbox for the author to self-review the code suggestions @@ -229,24 +233,9 @@ code_suggestions_self_review_text= "**Author self-review**: I have reviewed the approve_pr_on_self_review=false # if true, the PR will be auto-approved after the author clicks on the self-review checkbox fold_suggestions_on_self_review=true # if true, the code suggestions will be folded after the author clicks on the self-review checkbox -[pr_custom_prompt] # /custom_prompt # -prompt = """\ -The code suggestions should focus only on the following: -- ... -- ... -... -""" -suggestions_score_threshold=0 -num_code_suggestions_per_chunk=3 -self_reflect_on_custom_suggestions=true -enable_help_text=false - - [pr_add_docs] # /add_docs # extra_instructions = "" docs_style = "Sphinx" # "Google Style with Args, Returns, Attributes...etc", "Numpy Style", "Sphinx Style", "PEP257", "reStructuredText" -file = "" # in case there are several components with the same name, you can specify the relevant file -class_name = "" # in case there are several methods with the same name in the same file, you can specify the relevant class name [pr_update_changelog] # /update_changelog # push_changelog_changes=false @@ -254,28 +243,6 @@ extra_instructions = "" add_pr_link=true skip_ci_on_push=true -[pr_analyze] # /analyze # -enable_help_text=true - -[pr_test] # /test # -extra_instructions = "" -testing_framework = "" # specify the testing framework you want to use -num_tests=3 # number of tests to generate. max 5. -avoid_mocks=true # if true, the generated tests will prefer to use real objects instead of mocks -file = "" # in case there are several components with the same name, you can specify the relevant file -class_name = "" # in case there are several methods with the same name in the same file, you can specify the relevant class name -enable_help_text=false - -[pr_improve_component] # /improve_component # -num_code_suggestions=4 -extra_instructions = "" -file = "" # in case there are several components with the same name, you can specify the relevant file -class_name = "" # in case there are several methods with the same name in the same file, you can specify the relevant class name - -[pr_help] # /help # -force_local_db=false -num_retrieved_snippets=5 - [pr_config] # /config # [pr_help_docs] @@ -290,11 +257,16 @@ enable_help_text=false # The type of deployment to create. Valid values are 'app' or 'user'. deployment_type = "user" ratelimit_retries = 5 +# PyGithub 2.x enables request pacing and transient retries by default (0.25s +# between requests, 1s between writes, 10 retries); PyGithub 1.59 applied none. +# The defaults below mirror the 1.59 behaviour, so the 2.x upgrade stays neutral +# unless an operator opts in. Raise them to throttle or retry on the GitHub API. +seconds_between_requests = 0 # seconds between API requests; 0 = no pacing (1.59 behaviour) +seconds_between_writes = 0 # seconds between write calls; 0 = no pacing (1.59 behaviour) +api_retries = 0 # max retries per request with backoff; 0 = no retries (1.59 behaviour) polling_request_timeout = 10 # total seconds for comment-history fallback; positive values capped at 60 base_url = "https://api.github.com" -publish_inline_comments_fallback_with_verification = true try_fix_invalid_inline_comments = true -app_name = "pr-agent" ignore_bot_pr = true publish_as_check_run = false # when true, publish review/description/improve output as GitHub Checks instead of PR comments @@ -313,7 +285,6 @@ publish_as_check_run = false # when true, publish review/description/improve out [github_app] # these toggles allows running the github app from custom deployments -bot_user = "github-actions[bot]" override_deployment_type = true # settings for "pull_request" event handle_pr_actions = ['opened', 'reopened', 'ready_for_review'] @@ -322,11 +293,16 @@ feedback_on_draft_pr = false review_states = ["changes_requested"] review_author_types = ["User"] review_commands = [] +# Opt in to in-memory deduplication by X-GitHub-Delivery within each worker process. +# Active work stays protected; completed IDs expire after the deployment's push_trigger_pending_tasks_ttl +# (read at startup, 300 seconds by default). Failed or cancelled work can retry immediately. +# Repeated deliveries do not extend the TTL. Manual redeliveries are also suppressed during this window. +# State is cleared on restart and is not shared across workers or replicas. +webhook_delivery_deduplication = false # settings for "pull_request" event with "synchronize" action - used to detect and handle push triggers for new commits handle_push_trigger = false push_trigger_ignore_bot_commits = true push_trigger_ignore_merge_commits = true -push_trigger_wait_for_initial_review = true push_trigger_pending_tasks_backlog = true push_trigger_pending_tasks_ttl = 300 push_commands = [ @@ -399,9 +375,27 @@ webhook_password = "" # patch_server_token = "" [bitbucket_server] -# URL to the BitBucket Server instance # url = "https://git.bitbucket.com" -url = "" +url = "" # URL to the BitBucket Server instance + +[jira] +# Jira Cloud ticket context. Only Jira Cloud is supported; the base URL is built as +# https://.atlassian.net from the site name below, so repo configuration +# cannot redirect the authenticated request to another host. +# Credentials: fill and place in .secrets.toml, or set via environment variables +# (e.g. jira__jira_api_token). +# jira_site = "your-org" # the "" in https://.atlassian.net +# jira_api_email = "..." # Atlassian account email +# jira_api_token = "..." # Atlassian API token +# Custom field id holding acceptance criteria / requirements, mapped to the ticket +# "requirements" section. Instance-specific (e.g. "customfield_10127"); empty disables it. +jira_requirements_field = "" +# Optional allowlist of Jira project keys, e.g. ["PROJ", "OPS"]. When non-empty, key-shaped +# text with another prefix ("SHA-256", "UTF-8", "ISO-8601") is dropped before any lookup, +# so it no longer costs an authenticated 404 each. Entries must be plain upper-case keys; +# a supplied list with no valid entry disables the lookup rather than widening it. +# Empty (default) looks up every key found. +project_keys = [] [litellm] # use_client = false @@ -438,7 +432,7 @@ allow_fallbacks = true # when provider_order is set, allow routing beyond the li # Invalid reasoning_effort values are warned about and treated as unset. # Empty inherits config.reasoning_effort for models in SUPPORT_REASONING_EFFORT_MODELS. # Valid values: "none", "minimal", "low", "medium", "high", "xhigh", "max". -# OpenRouter normalizes "max" to "xhigh" to match LiteLLM 1.98.0. +# OpenRouter normalizes "max" to "xhigh" for LiteLLM/OpenRouter compatibility. # Model-specific support varies; mandatory reasoning models reject "none". reasoning_effort = "" # A positive value overrides global effort and non-none OpenRouter-specific efforts. @@ -462,11 +456,12 @@ rules = [] # e.g. [{ max_hunks = 3, model = "gpt-5.6-luna" }, { max_hunks = 15, [otel] # Keys map to settings.get("OTEL.*") via section-key dot notation. is_enabled = false # disabled by default; set to true to enable telemetry -exporter_type = "console" # "console", "otlp", or "none" +exporter_type = "console" # "console", "otlp", "prometheus", or "none" service_name = "pr-agent" environment = "development" # "development", "staging", "production", etc. otlp_timeout = 3 # seconds; hard deadline per OTLP export call (incl. retries). Bounds CLI-exit/request stalls when the collector is unreachable; batches slower than this are dropped. otlp_protocol = "http" # "http" (default; exporter ships with pr-agent) or "grpc" (requires the otel-grpc extra: pip install pr-agent[otel-grpc]) +prometheus_multiproc_dir = "/tmp/pr-agent-prometheus" # shared dir for the prometheus exporter's per-worker state files; required when exporter_type = "prometheus" (multiprocess gunicorn deployments); created 0700 and must stay owned by the run user (symlink or foreign-owner paths are refused) include_pr_url = false # set to true to attach PR URLs to spans (may expose private repo names) include_error_details = false # set to true to attach exception messages and rejected-command text to spans (may expose PR URLs, repo names, or other request content) @@ -474,20 +469,15 @@ include_error_details = false # set to true to attach exception messages and rej skip_comments = false force_update_dataset = false max_issues_to_scan = 500 -vectordb = "pinecone" # options: "pinecone", "lancedb", "qdrant" - -[pr_find_similar_component] -class_name = "" -file = "" -search_from_org = false -allow_fallback_less_words = true -number_of_keywords = 5 -number_of_results = 5 +vectordb = "lancedb" # options: "pinecone", "lancedb", "qdrant" [pinecone] # fill and place in .secrets.toml #api_key = ... -# environment = "gcp-starter" +# Serverless index deployment for the modern SDK. `cloud` is one of "aws", +# "gcp" or "azure"; pick a `region` offered by that cloud. +cloud = "aws" +region = "us-east-1" [lancedb] uri = "./lancedb" @@ -511,19 +501,6 @@ enabled = false paths = [] # directories to scan recursively for "*/SKILL.md"; supports ~ and $VAR max_skills_tokens = 8000 # token budget for the combined skills_context block -[best_practices] -content = "" -organization_name = "" -max_lines_allowed = 800 -enable_global_best_practices = false - -[auto_best_practices] -enable_auto_best_practices = true # public - general flag to disable all auto best practices usage -utilize_auto_best_practices = true # public - disable usage of auto best practices in the 'improve' tool -extra_instructions = "" # public - extra instructions to the auto best practices generation prompt -content = "" -max_patterns = 5 # max number of patterns to be detected - [artifacts] # Enable artifact injection into tool prompts (off by default; auto-enabled when artifact_path input is set) enable = false @@ -531,9 +508,9 @@ enable = false artifact_path = "" # Custom instructions appended after the artifact content (leave empty for a sensible default) artifact_instructions = "" -# Label shown to the AI — defaults to the filename when empty +# Label shown to the AI — defaults to the filename when empty. artifact_label = "" -# Which tools receive artifact context +# Which tools receive artifact context. target_tools = ["pr_reviewer", "pr_description", "pr_code_suggestions"] # Max artifact size in characters (content is truncated if exceeded) max_artifact_size = 50000 diff --git a/pr_agent/settings/pr_add_docs.toml b/pr_agent/settings/pr_add_docs.toml index 8efb8b945b..e9fdc41e43 100644 --- a/pr_agent/settings/pr_add_docs.toml +++ b/pr_agent/settings/pr_add_docs.toml @@ -56,7 +56,7 @@ Example output: Code Documentation: - relevant file: |- src/file1.py - relevant lines: 12 + relevant line: 12 doc placement: after documentation: |- \"\"\" diff --git a/pr_agent/settings/pr_evaluate_prompt_response.toml b/pr_agent/settings/pr_evaluate_prompt_response.toml index 9a68139307..b6f821befa 100644 --- a/pr_agent/settings/pr_evaluate_prompt_response.toml +++ b/pr_agent/settings/pr_evaluate_prompt_response.toml @@ -44,9 +44,9 @@ After that, rank each response. Criterions to rank each response: - Don't necessarily rank higher a response that is longer. A shorter response might be better if it is more concise, and still addresses the task better. -The output must be a YAML object equivalent to type $PRRankRespones, according to the following Pydantic definitions: +The output must be a YAML object equivalent to type $PRRankResponses, according to the following Pydantic definitions: ===== -class PRRankRespones(BaseModel): +class PRRankResponses(BaseModel): which_response_was_better: Literal[0, 1, 2] = Field(description="A number indicating which response was better. 0 means both responses are equally good.") why: str = Field(description="In a short and concise manner, explain why the chosen response is better than the other. Be specific and give examples if relevant.") score_response1: int = Field(description="A score between 1 and 10, indicating the quality of the response1, based on the criterions mentioned in the prompt.") @@ -56,8 +56,8 @@ class PRRankRespones(BaseModel): Example output: ```yaml -which_response_was_better: "X" -why: "Response X is better because it is more practical, and addresses the task requirements better since ..." +which_response_was_better: 1 +why: "Response 1 is better because it is more practical, and addresses the task requirements better since ..." score_response1: ... score_response2: ... ``` diff --git a/pr_agent/settings/pr_help_prompts.toml b/pr_agent/settings/pr_help_prompts.toml index 274940fd2a..ade4c5b667 100644 --- a/pr_agent/settings/pr_help_prompts.toml +++ b/pr_agent/settings/pr_help_prompts.toml @@ -1,5 +1,5 @@ [pr_help_prompts] -system="""You are Doc-helper, a language model designed to answer questions about a documentation website for an open-source project called "PR-Agent" (recently renamed to "Qodo Merge"). +system="""You are Doc-helper, a language model designed to answer questions about the open-source PR-Agent project. You will receive a question, and the full documentation website content. Your goal is to provide the best answer to the question using the documentation provided. diff --git a/pr_agent/settings/pr_reviewer_prompts.toml b/pr_agent/settings/pr_reviewer_prompts.toml index e797cc4536..83f17866e1 100644 --- a/pr_agent/settings/pr_reviewer_prompts.toml +++ b/pr_agent/settings/pr_reviewer_prompts.toml @@ -100,10 +100,10 @@ class Review(BaseModel): estimated_effort_to_review_[1-5]: int = Field(description="Estimate, on a scale of 1-5 (inclusive), the time and effort required to review this PR by an experienced and knowledgeable developer. 1 means short and easy review, 5 means long and hard review. Take into account the size, complexity, quality, and the needed changes of the PR code diff.") {%- endif %} {%- if require_risk_assessment %} - risk_level: str = Field(description="Overall risk level of this PR. Answer with exactly one of: low, medium, high. Use high only when the PR introduces a clear bug, security concern, or major logic risk. Use medium when the PR is not clearly broken but contains non-trivial areas that require careful human verification. Use low when the PR is small, low-impact, and no important issues are identified.") + risk_level: Literal["low", "medium", "high"] = Field(description="Overall risk level of this PR. Answer with exactly one of: low, medium, high. Use high only when the PR introduces a clear bug, security concern, or major logic risk. Use medium when the PR is not clearly broken but contains non-trivial areas that require careful human verification. Use low when the PR is small, low-impact, and no important issues are identified.") {%- endif %} {%- if require_merge_recommendation %} - merge_recommendation: str = Field(description="Overall merge recommendation for this PR. Answer with exactly one of: safe_to_merge, merge_with_caution, changes_required. Use changes_required when there are clear issues that should be fixed before merge. Use merge_with_caution when the PR seems acceptable but still deserves focused reviewer attention. Use safe_to_merge when no important blockers or risks are identified.") + merge_recommendation: Literal["safe_to_merge", "merge_with_caution", "changes_required"] = Field(description="Overall merge recommendation for this PR. Answer with exactly one of: safe_to_merge, merge_with_caution, changes_required. Use changes_required when there are clear issues that should be fixed before merge. Use merge_with_caution when the PR seems acceptable but still deserves focused reviewer attention. Use safe_to_merge when no important blockers or risks are identified.") {%- endif %} {%- if require_priority_files %} review_priority_files: List[str] = Field(description="A short list of the most important files a human reviewer should inspect first. Return an empty list if the PR is too small or no file deserves special attention.") @@ -112,10 +112,10 @@ class Review(BaseModel): contribution_time_cost_estimate: ContributionTimeCostEstimate = Field(description="An estimate of the time required to implement the changes, based on the quantity, quality, and complexity of the contribution, as well as the context from the PR description and commit messages.") {%- endif %} {%- if require_score %} - score: str = Field(description="Rate this PR on a scale of 0-100 (inclusive), where 0 means the worst possible PR code, and 100 means PR code of the highest quality, without any bugs or performance issues, that is ready to be merged immediately and run in production at scale.") + score: int = Field(description="Rate this PR on a scale of 0-100 (inclusive), where 0 means the worst possible PR code, and 100 means PR code of the highest quality, without any bugs or performance issues, that is ready to be merged immediately and run in production at scale.") {%- endif %} {%- if require_tests %} - relevant_tests: str = Field(description="yes/no question: does this PR have relevant tests added or updated?") + relevant_tests: Literal["Yes", "No"] = Field(description="Does this PR have relevant tests added or updated? Answer exactly Yes or No.") {%- endif %} {%- if question_str %} insights_from_user_answers: str = Field(description="shortly summarize the insights you gained from the user's answers to the questions") @@ -149,20 +149,17 @@ review: ... not_compliant_requirements: | ... - overall_compliance_level: | + requires_further_human_verification: | ... {%- endif %} {%- if require_estimate_effort_to_review %} - estimated_effort_to_review_[1-5]: | - 3 + estimated_effort_to_review_[1-5]: 3 {%- endif %} {%- if require_risk_assessment %} - risk_level: | - low + risk_level: low {%- endif %} {%- if require_merge_recommendation %} - merge_recommendation: | - safe_to_merge + merge_recommendation: safe_to_merge {%- endif %} {%- if require_priority_files %} review_priority_files: @@ -300,12 +297,11 @@ review: ... not_compliant_requirements: | ... - overall_compliance_level: | + requires_further_human_verification: | ... {%- endif %} {%- if require_estimate_effort_to_review %} - estimated_effort_to_review_[1-5]: | - 3 + estimated_effort_to_review_[1-5]: 3 {%- endif %} {%- if require_risk_assessment %} risk_level: | diff --git a/pr_agent/telemetry/config.py b/pr_agent/telemetry/config.py index a192f17256..d320784dda 100644 --- a/pr_agent/telemetry/config.py +++ b/pr_agent/telemetry/config.py @@ -1,9 +1,10 @@ from pr_agent.algo.utils import get_version from pr_agent.config_loader import get_settings from pr_agent.log import get_logger +from pr_agent.telemetry.prometheus_multiproc import ensure_prometheus_multiproc_dir from pr_agent.telemetry.types import ExporterType, OtlpProtocol, TelemetryConfig -VALID_EXPORTER_TYPES = {ExporterType.CONSOLE, ExporterType.OTLP, ExporterType.NONE} +VALID_EXPORTER_TYPES = {ExporterType.CONSOLE, ExporterType.OTLP, ExporterType.PROMETHEUS, ExporterType.NONE} VALID_OTLP_PROTOCOLS = {OtlpProtocol.HTTP, OtlpProtocol.GRPC} @@ -42,7 +43,8 @@ def get_otel_config() -> TelemetryConfig: return TelemetryConfig() # Fail closed: opted-in telemetry content was consented for the OTLP destination - # only — never redirect it to another exporter (console = process logs). + # only — never redirect it to another exporter (console = process logs; + # prometheus = a local scrape endpoint the operator explicitly selected). if exporter_type == ExporterType.OTLP and not otlp_endpoint: get_logger().warning( "OTEL.EXPORTER_TYPE is 'otlp' but OTEL.OTLP_ENDPOINT is not configured. " @@ -50,6 +52,15 @@ def get_otel_config() -> TelemetryConfig: ) return TelemetryConfig() + # A multiprocess scrape must write per-pid state files into a shared dir. + # The env var has to be set before prometheus_client is imported in any + # worker (gunicorn master also provisions it in `when_ready`). + prometheus_multiproc_dir = None + if exporter_type == ExporterType.PROMETHEUS: + prometheus_multiproc_dir = ensure_prometheus_multiproc_dir( + str(settings.get("OTEL.PROMETHEUS_MULTIPROC_DIR", "/tmp/pr-agent-prometheus")) + ) + return TelemetryConfig( is_enabled=True, exporter_type=exporter_type, @@ -59,7 +70,8 @@ def get_otel_config() -> TelemetryConfig: otlp_endpoint=otlp_endpoint, otlp_headers=_parse_otlp_headers(otlp_headers_raw) if otlp_headers_raw else None, otlp_timeout=int(settings.get("OTEL.OTLP_TIMEOUT", 3)), - otlp_protocol=otlp_protocol + otlp_protocol=otlp_protocol, + prometheus_multiproc_dir=prometheus_multiproc_dir, ) diff --git a/pr_agent/telemetry/meter.py b/pr_agent/telemetry/meter.py index 908585b630..8d6139003c 100644 --- a/pr_agent/telemetry/meter.py +++ b/pr_agent/telemetry/meter.py @@ -73,6 +73,12 @@ def _create_metric_exporter(config): if config.otlp_headers: kwargs["headers"] = config.otlp_headers return OTLPMetricExporter(**kwargs) + elif config.exporter_type == ExporterType.PROMETHEUS: + # Imported lazily so the gunicorn master (preload_app) never pulls in + # prometheus_client before the multiprocess dir is provisioned. + from pr_agent.telemetry.prometheus import PrometheusMetricExporter + + return PrometheusMetricExporter() return None diff --git a/pr_agent/telemetry/prometheus.py b/pr_agent/telemetry/prometheus.py new file mode 100644 index 0000000000..8d62454897 --- /dev/null +++ b/pr_agent/telemetry/prometheus.py @@ -0,0 +1,209 @@ +"""Native Prometheus scrape endpoint for the OpenTelemetry command telemetry. + +When ``OTEL.EXPORTER_TYPE = "prometheus"`` the ``pr_agent.commands`` counter +stays an OpenTelemetry instrument; this module translates the SDK's aggregated +data points into ``prometheus_client`` metric objects instead of shipping them +to a collector. + +Gunicorn runs with ``preload_app = True`` (fork model), so the translation +happens on the worker side: each worker's values land in its own +multiprocess state file under ``PROMETHEUS_MULTIPROC_DIR``, and ``/metrics`` +serves a ``MultiProcessCollector`` that merges every worker's file. + +prometheus_client decides at import time whether metrics are multiprocess-aware +(from ``PROMETHEUS_MULTIPROC_DIR``), and gunicorn's master imports the app +before ``when_ready`` provisions that directory. This module therefore never +imports prometheus_client at module level: everything touching it is deferred +to the worker or to a scrape request, so a master process merely importing the +app can neither create metrics in the wrong mode nor see a stale directory. + +``opentelemetry-exporter-prometheus`` is deliberately not a dependency: its +``PrometheusMetricReader`` documents "No multiprocessing support", which is the +requirement this module exists for, and it ships as a 0.x pre-release. +""" + +import re +from typing import Any, Dict, Tuple + +from opentelemetry.sdk.metrics.export import ( + Gauge as OTelGauge, +) +from opentelemetry.sdk.metrics.export import ( + Histogram as OTelHistogram, +) +from opentelemetry.sdk.metrics.export import ( + MetricExporter, + MetricExportResult, + Sum, +) +from starlette.responses import Response + +from pr_agent.config_loader import get_settings +from pr_agent.log import get_logger +from pr_agent.telemetry.types import ExporterType + +_NAME_INVALID_CHARS = re.compile(r"[^a-zA-Z0-9_:]") + +_prometheus_client = None +_prometheus_registry = None + + +def _client() -> Any: + """Import prometheus_client on first use (never during app import in the master).""" + global _prometheus_client + if _prometheus_client is None: + import prometheus_client + + # The multiprocess dir/env is provisioned by config/`when_ready` before + # any worker reaches this point, so the client picks the right mode here. + _prometheus_client = prometheus_client + return _prometheus_client + + +def prometheus_metrics_enabled() -> bool: + """True when the [otel] prometheus exporter is selected and telemetry is on.""" + settings = get_settings() + return bool(settings.get("OTEL.IS_ENABLED", False)) and ( + settings.get("OTEL.EXPORTER_TYPE") == ExporterType.PROMETHEUS + ) + + +def prometheus_registry(): + """The registry metric objects are created on (single-process view).""" + return _registry() + + +def _registry(): + global _prometheus_registry + if _prometheus_registry is None: + _prometheus_registry = _client().CollectorRegistry(auto_describe=False) + return _prometheus_registry + + +def prometheus_response() -> Response: + """Render ``/metrics`` for this process. + + Uses a ``MultiProcessCollector`` when prometheus_client imported in + multiprocess mode (workers under gunicorn with PROMETHEUS_MULTIPROC_DIR + set); otherwise serves this process's own registry (single-process + deployments such as plain `uvicorn`). + """ + client = _client() + # prometheus_client decides its storage mode when `values` is first imported + # (from PROMETHEUS_MULTIPROC_DIR); MultiProcessValue is chosen per worker. + import prometheus_client.multiprocess + import prometheus_client.values as prometheus_values + + if prometheus_values.ValueClass._multiprocess: + registry = client.CollectorRegistry() + prometheus_client.multiprocess.MultiProcessCollector(registry) + else: + registry = _registry() + return Response( + content=client.generate_latest(registry), media_type=client.CONTENT_TYPE_LATEST + ) + + +def attach_metrics_endpoint(router) -> None: + """Attach a lazy ``GET /metrics`` route to a FastAPI router. + + The route raises no prometheus_client import at registration time — the + gunicorn master (which imports the app under ``preload_app``) must never + import it before ``when_ready`` provisions the multiprocess directory. The + heavy import happens on the first request, inside the serving worker. + """ + + def _metrics_endpoint(): + return prometheus_response() + + router.add_api_route( + "/metrics", _metrics_endpoint, methods=["GET"], include_in_schema=False + ) + + +class PrometheusMetricExporter(MetricExporter): + """Translate SDK `MetricsData` into prometheus_client metrics. + + Designed for use with ``PeriodicExportingMetricReader`` (DELTA temporality + for counters, accumulating per worker). Increments land in the worker's + multiprocess state file when prometheus_client imported in multiprocess mode. + """ + + def __init__(self): + super().__init__() + # metric name -> (label_names, prometheus_client metric object) + self._counters: Dict[str, Tuple] = {} + self._gauges: Dict[str, Tuple] = {} + + def export(self, metrics_data, timeout_millis=10_000, **kwargs) -> MetricExportResult: + try: + for resource_metric in metrics_data.resource_metrics: + for scope_metric in resource_metric.scope_metrics: + for metric in scope_metric.metrics: + self._translate(metric) + except Exception as e: + get_logger().warning(f"Failed to export metrics to Prometheus: {e}") + return MetricExportResult.FAILURE + return MetricExportResult.SUCCESS + + def force_flush(self, timeout_millis=10_000) -> bool: + return True + + def shutdown(self, timeout_millis=30_000, **kwargs) -> None: + pass + + def _translate(self, metric) -> None: + name = _sanitize_metric_name(metric.name) + help_text = metric.description or metric.name + if isinstance(metric.data, Sum): + if metric.data.is_monotonic: + for point in metric.data.data_points: + obj, labels = self._metric(self._counters, "Counter", name, help_text, point) + obj.labels(**labels).inc(point.value) + else: + for point in metric.data.data_points: + obj, labels = self._metric(self._gauges, "Gauge", name, help_text, point) + obj.labels(**labels).set(point.value) + elif isinstance(metric.data, OTelGauge): + for point in metric.data.data_points: + obj, labels = self._metric(self._gauges, "Gauge", name, help_text, point) + obj.labels(**labels).set(point.value) + elif isinstance(metric.data, OTelHistogram): + get_logger().warning( + f"Skipping histogram metric '{metric.name}': Prometheus histograms are " + f"not exported yet (no command telemetry uses one today)." + ) + else: + get_logger().warning( + f"Skipping unsupported metric '{metric.name}' of type {type(metric.data)}." + ) + + def _metric(self, cache, metric_class_name, name, help_text, point): + labels = _point_labels(point) + entry = cache.get(name) + if entry is None: + client = _client() + metric_class = getattr(client, metric_class_name) + label_names = tuple(sorted(labels)) + metric_object = metric_class( + name, help_text, labelnames=label_names, registry=_registry() + ) + entry = cache[name] = (label_names, metric_object) + label_names, metric_object = entry + # A family's label set is fixed on the first data point; drop later + # attributes we cannot attach and back-fill missing ones with "". + normalized = {label: labels.get(label, "") for label in label_names} + return metric_object, normalized + + +def _sanitize_metric_name(name: str) -> str: + return _NAME_INVALID_CHARS.sub("_", name) + + +def _sanitize_label_key(key: str) -> str: + key = _NAME_INVALID_CHARS.sub("_", key) + return "_" + key if key[0].isdigit() else key + + +def _point_labels(point) -> dict: + return {_sanitize_label_key(key): str(value) for key, value in point.attributes.items()} diff --git a/pr_agent/telemetry/prometheus_multiproc.py b/pr_agent/telemetry/prometheus_multiproc.py new file mode 100644 index 0000000000..1b767368ee --- /dev/null +++ b/pr_agent/telemetry/prometheus_multiproc.py @@ -0,0 +1,53 @@ +"""Prometheus multiprocess state-directory setup. + +`prometheus_client` decides whether metrics are multiprocess-capable when its +`metrics` module is first imported, by looking at the `PROMETHEUS_MULTIPROC_DIR` +environment variable. Under gunicorn's fork model (`preload_app = True`) the +state directory must therefore exist and the variable be set *before* any +worker imports prometheus_client. + +This module imports nothing but the standard library so it is safe to call from +the gunicorn master (via `gunicorn_config.when_ready`) and from +`pr_agent.telemetry.config`, both before any worker has imported +prometheus_client. +""" + +import os + +PROMETHEUS_MULTIPROC_DIR_ENV = "PROMETHEUS_MULTIPROC_DIR" + + +def ensure_prometheus_multiproc_dir(path: str) -> str: + """Set ``PROMETHEUS_MULTIPROC_DIR`` and create or validate the directory. + + Returns the normalized path. Idempotent; safe to call from the gunicorn + master and from each worker's first telemetry init. + + The state files hold per-pid metric data that any local user with read + access could observe, so a predictable default inside a world-writable + prefix like /tmp must not be redirected or shared: a symlink is refused, + the directory is created (and tightened to) 0700, and a path owned by + another user is rejected instead of followed. + """ + env = PROMETHEUS_MULTIPROC_DIR_ENV + if path: + os.environ[env] = path + target = os.environ[env] + if os.path.islink(target): + raise RuntimeError( + f"Refusing to use a symlinked prometheus multiprocess state directory: {target!r}" + ) + os.makedirs(target, mode=0o700, exist_ok=True) + st = os.stat(target) + if st.st_uid != os.getuid(): + raise RuntimeError( + f"Prometheus multiprocess state directory {target!r} is owned by uid {st.st_uid}, " + f"not the current user's uid {os.getuid()}" + ) + os.chmod(target, 0o700) + return target + + +def prometheus_multiproc_dir() -> str | None: + """The configured multiprocess state directory, if any.""" + return os.environ.get(PROMETHEUS_MULTIPROC_DIR_ENV) diff --git a/pr_agent/telemetry/tracer.py b/pr_agent/telemetry/tracer.py index 7a70fbcd31..043632b384 100644 --- a/pr_agent/telemetry/tracer.py +++ b/pr_agent/telemetry/tracer.py @@ -72,4 +72,9 @@ def _create_exporter(config): if config.otlp_headers: kwargs['headers'] = config.otlp_headers return OTLPSpanExporter(**kwargs) + elif config.exporter_type == ExporterType.PROMETHEUS: + get_logger().warning( + "OTEL.EXPORTER_TYPE is 'prometheus' (metrics only); spans are not exported." + ) + return None return None diff --git a/pr_agent/telemetry/types.py b/pr_agent/telemetry/types.py index f104cf70d9..f7e14e9129 100644 --- a/pr_agent/telemetry/types.py +++ b/pr_agent/telemetry/types.py @@ -6,6 +6,7 @@ class ExporterType: """Canonical exporter type values for OTEL.EXPORTER_TYPE.""" OTLP = "otlp" CONSOLE = "console" + PROMETHEUS = "prometheus" NONE = "none" @@ -26,3 +27,4 @@ class TelemetryConfig: otlp_headers: Optional[Dict[str, str]] = None otlp_timeout: int = 3 # seconds; hard deadline per export call, retries included otlp_protocol: str = OtlpProtocol.HTTP + prometheus_multiproc_dir: Optional[str] = None # shared dir for multiprocess metrics diff --git a/pr_agent/tools/pr_code_suggestions.py b/pr_agent/tools/pr_code_suggestions.py index 38726ae31f..ed068f9101 100644 --- a/pr_agent/tools/pr_code_suggestions.py +++ b/pr_agent/tools/pr_code_suggestions.py @@ -1,6 +1,7 @@ import asyncio import copy import difflib +import math import re import textwrap import traceback @@ -10,20 +11,21 @@ from jinja2 import Environment, StrictUndefined -from pr_agent.algo import MAX_TOKENS from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler from pr_agent.algo.git_patch_processing import decouple_and_convert_to_hunks_with_lines_numbers from pr_agent.algo.pr_processing import ( + OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD, _get_all_models, add_ai_metadata_to_diff_files, + get_effective_fallback_chain, get_pr_diff, get_pr_multi_diffs, retry_with_fallback_models, ) from pr_agent.algo.prompt_fragments import render_diff_hunk_format from pr_agent.algo.repo_context import build_repo_context -from pr_agent.algo.run_details import init_run_details +from pr_agent.algo.run_details import init_run_details, record_model_used from pr_agent.algo.skills_loader import get_skills_context from pr_agent.algo.token_handler import TokenHandler from pr_agent.algo.utils import ( @@ -43,10 +45,7 @@ show_run_details, ) from pr_agent.config_loader import get_settings, get_verbosity_level -from pr_agent.git_providers import ( - GithubProvider, - get_git_provider_with_context, -) +from pr_agent.git_providers import get_git_provider_with_context from pr_agent.git_providers.git_provider import GitProvider, IncrementalPR, get_main_pr_language from pr_agent.log import get_logger from pr_agent.servers.help import HelpMessage @@ -111,6 +110,19 @@ def _supports_code_suggestion_state(git_provider) -> bool: return callable(supports) and bool(supports()) +def _supports_persistent_progress_comment(git_provider) -> bool: + """Whether a published progress comment can later be edited in place or removed. + + Hosted providers turn the progress note into the final suggestions comment (edit it) or + delete it on failure or cancellation, so publishing it up front is worthwhile. Output-only + providers (plain-diff) write every non-temporary comment straight through to their output + and can do neither; persisting the progress there would leak a stale "Preparing + suggestions..." document ahead of the final result. + """ + return (git_provider.is_supported("edit_comment") + and git_provider.is_supported("remove_comment")) + + def _edit_comment_safely(git_provider, comment, body: str) -> bool: try: result = git_provider.edit_comment(comment, body) @@ -284,9 +296,15 @@ async def run(self): not get_settings().config.get('is_auto_command', False)): if self.git_provider.is_supported("gfm_markdown"): # The progress comment later becomes the final suggestions comment (edited in place), - # so it must already be a thread when threaded output is requested. - self.progress_response = self.git_provider.publish_comment(self.progress, - **self._improve_thread_kwargs()) + # so it must already be a thread when threaded output is requested. Output-only + # providers (plain-diff) cannot edit or remove it afterwards, so for those the + # progress is a temporary placeholder that is never persisted. + if _supports_persistent_progress_comment(self.git_provider): + self.progress_response = self.git_provider.publish_comment(self.progress, + **self._improve_thread_kwargs()) + else: + self.progress_response = self.git_provider.publish_comment( + self.progress, is_temporary=True, **self._improve_thread_kwargs()) else: self.progress_response = self.git_provider.publish_comment( "Preparing suggestions...", is_temporary=True) @@ -318,6 +336,20 @@ async def run(self): if ((not get_settings().pr_code_suggestions.commitable_code_suggestions) and self.git_provider.is_supported("gfm_markdown")): + # Drop suggestions that can't be anchored in the diff (unresolved + # sentinels, zero/negative or reversed line ranges, or positive + # ranges that fall outside the changed lines of the relevant file) + # up front; when nothing survives, route the outcome through + # publish_no_suggestions() so it honors publish_output_no_suggestions + # and emits the accurate coverage footer instead of a header-only table. + data['code_suggestions'] = [ + suggestion for suggestion in data['code_suggestions'] + if self._is_suggestion_line_range_valid(suggestion) + ] + if not data['code_suggestions']: + await self.publish_no_suggestions() + return + # generate summarized suggestions pr_body = self.generate_summarized_suggestions(data) pr_body += self._get_suggestions_coverage_footer() @@ -329,7 +361,7 @@ async def run(self): # add usage guide if (get_settings().pr_code_suggestions.enable_chat_text and get_settings().config.is_auto_command - and isinstance(self.git_provider, GithubProvider)): + and self.git_provider.supports_pr_chat()): pr_body += "\n\n>💡 Need additional feedback ? start a [PR chat](https://chromewebstore.google.com/detail/ephlnjeghhogofkifjloamocljapahnl) \n\n" if get_settings().pr_code_suggestions.enable_help_text: pr_body += "
      \n\n
      💡 Tool usage guide:
      \n\n" @@ -787,15 +819,6 @@ def _with_identity(comment_text: str) -> str: new_comment = git_provider.publish_comment(pr_comment, **({"as_thread": True} if as_thread else {})) return new_comment - def extract_link(self, s): - r = re.compile(r"") - match = r.search(s) - - up_to_commit_txt = "" - if match: - up_to_commit_txt = f" up to commit {match.group(0)[4:-3].strip()}" - return up_to_commit_txt - async def _prepare_prediction(self, model: str) -> dict: self.patches_diff = get_pr_diff(self.git_provider, self.token_handler, @@ -815,13 +838,17 @@ async def _prepare_prediction(self, model: str) -> dict: data = self.prediction return data - async def _get_prediction(self, model: str, patches_diff: str, patches_diff_no_line_number: str) -> dict: + def _render_prediction_prompts(self, patches_diff: str, patches_diff_no_line_number: str) -> tuple[str, str]: variables = copy.deepcopy(self.vars) variables["diff"] = patches_diff # update diff variables["diff_no_line_numbers"] = patches_diff_no_line_number # update diff environment = Environment(undefined=StrictUndefined) system_prompt = environment.from_string(self.pr_code_suggestions_prompt_system).render(variables) user_prompt = environment.from_string(self.pr_code_suggestions_prompt_user).render(variables) + return system_prompt, user_prompt + + async def _get_prediction(self, model: str, patches_diff: str, patches_diff_no_line_number: str) -> dict: + system_prompt, user_prompt = self._render_prediction_prompts(patches_diff, patches_diff_no_line_number) response, finish_reason = await self.ai_handler.chat_completion( model=model, temperature=get_settings().config.temperature, system=system_prompt, user=user_prompt) if not get_settings().config.publish_output: @@ -836,8 +863,8 @@ async def _get_prediction(self, model: str, patches_diff: str, patches_diff_no_l if response_reflect: await self.analyze_self_reflection_response(data, response_reflect) else: - # get_logger().error(f"Could not self-reflect on suggestions. using default score 7") - for i, suggestion in enumerate(data["code_suggestions"]): + get_logger().warning("Could not self-reflect on suggestions; using default score 7") + for suggestion in data["code_suggestions"]: suggestion["score"] = 7 suggestion["score_why"] = "" @@ -880,7 +907,18 @@ async def _self_reflect_with_fallback(self, suggestion_list: List, patches_diff: async def analyze_self_reflection_response(self, data, response_reflect): response_reflect_yaml = load_yaml(response_reflect) + if not isinstance(response_reflect_yaml, dict): + get_logger().warning( + "Self-reflection feedback was not a mapping; line anchors will not be resolved" + ) + return code_suggestions_feedback = response_reflect_yaml.get("code_suggestions", []) + if not isinstance(code_suggestions_feedback, list): + get_logger().warning( + "Self-reflection feedback 'code_suggestions' was not a list; " + "line anchors will not be resolved" + ) + return if code_suggestions_feedback and len(code_suggestions_feedback) == len(data["code_suggestions"]): for i, suggestion in enumerate(data["code_suggestions"]): try: @@ -931,6 +969,11 @@ async def analyze_self_reflection_response(self, data, response_reflect): suggestion['existing_code'] = "" except Exception as e: get_logger().error(f"Error processing suggestion {i + 1}, error: {e}") + else: + get_logger().warning( + f"Self-reflection feedback covered {len(code_suggestions_feedback)} suggestion(s) instead of " + f"{len(data['code_suggestions'])}; line anchors will not be resolved" + ) @staticmethod def _truncate_if_needed(suggestion): @@ -946,6 +989,81 @@ def _truncate_if_needed(suggestion): suggestion['_is_truncated'] = True return suggestion + @staticmethod + def _parse_line_number(value) -> Optional[int]: + """Convert an anchor value to an int, or None when it cannot be a line number. + + Rejects booleans, fractional floats (int() truncates them), and non-finite + floats (int() raises OverflowError) instead of silently coercing them. + """ + if isinstance(value, bool): + return None + if isinstance(value, float) and (not math.isfinite(value) or not value.is_integer()): + return None + try: + return int(value) + except (TypeError, ValueError, OverflowError): + return None + + def _is_suggestion_line_range_valid(self, suggestion: dict) -> bool: + relevant_lines_start = self._parse_line_number(suggestion.get('relevant_lines_start')) + relevant_lines_end = self._parse_line_number(suggestion.get('relevant_lines_end')) + relevant_file = suggestion.get('relevant_file') + if relevant_lines_start is None or relevant_lines_end is None: + get_logger().warning("Skipping a suggestion without a valid line range", + artifact={'relevant_file': relevant_file, + 'one_sentence_summary': suggestion.get('one_sentence_summary')}) + return False + if relevant_lines_start < 1 or relevant_lines_end < relevant_lines_start: + get_logger().warning("Skipping a suggestion with an invalid line range", + artifact={'relevant_file': relevant_file, + 'one_sentence_summary': suggestion.get('one_sentence_summary'), + 'relevant_lines_start': relevant_lines_start, + 'relevant_lines_end': relevant_lines_end}) + return False + suggestion['relevant_lines_start'] = relevant_lines_start + suggestion['relevant_lines_end'] = relevant_lines_end + if not self._is_suggestion_line_range_in_diff( + relevant_file, relevant_lines_start, relevant_lines_end): + return False + return True + + def _is_suggestion_line_range_in_diff( + self, + relevant_file, + relevant_lines_start: int, + relevant_lines_end: int) -> bool: + """Reject suggestions whose range cannot be resolved in the relevant file. + + The model reflects on the extended patch, so accept context outside the + raw hunk when the complete head file contains it, as in _validate_suggestion. + Without complete file content, use the raw hunk as the reference. + """ + if not isinstance(relevant_file, str) or not relevant_file.strip(): + get_logger().warning("Skipping a suggestion whose file is missing", + artifact={'relevant_file': relevant_file, + 'one_sentence_summary': None}) + return False + diff_file = self._get_diff_file(relevant_file.strip()) + if diff_file is None: + get_logger().warning("Skipping a suggestion whose file is not part of the PR diff", + artifact={'relevant_file': relevant_file.strip(), + 'relevant_lines_start': relevant_lines_start, + 'relevant_lines_end': relevant_lines_end}) + return False + if diff_file.head_file and getattr(diff_file, "head_file_is_complete", True): + range_resolvable = relevant_lines_end <= len(diff_file.head_file.splitlines()) + else: + range_resolvable = self._get_patch_range_lines( + diff_file.patch, relevant_lines_start, relevant_lines_end) is not None + if not range_resolvable: + get_logger().warning("Skipping a suggestion whose line range is not within the file", + artifact={'relevant_file': relevant_file.strip(), + 'relevant_lines_start': relevant_lines_start, + 'relevant_lines_end': relevant_lines_end}) + return False + return True + def _prepare_pr_code_suggestions(self, predictions: str) -> Dict: data = load_yaml(predictions.strip(), keys_fix_yaml=["relevant_file", "suggestion_content", "existing_code", "improved_code"], @@ -1055,7 +1173,9 @@ def _limit_suggestions_per_file(self, suggestions: List[Dict]) -> List[Dict]: async def push_inline_code_suggestions(self, data, include_coverage_footer: bool = True) -> None: code_suggestions = [] + artifact_suggestions = [] fallback_comments = [] + artifact_batch_published = False coverage_footer = self._get_suggestions_coverage_footer() if include_coverage_footer else "" supports_suggestions_artifact = self.git_provider.supports_code_suggestions_artifact() is True @@ -1126,20 +1246,26 @@ async def push_inline_code_suggestions(self, data, include_coverage_footer: bool elif requires_pr_fallback: body += f"\n\nNot offered as a committable change because {fallback_reason}." - # Keep safety-rejected suggestions out of provider patch APIs while preserving standalone artifacts. + rendered_suggestion = {'body': body, 'relevant_file': relevant_file, + 'relevant_lines_start': relevant_lines_start, + 'relevant_lines_end': relevant_lines_end, + 'original_suggestion': d} + if supports_suggestions_artifact: + artifact_suggestions.append(rendered_suggestion) + + # Keep safety-rejected suggestions out of provider patch APIs and retain fallback recovery text. if not has_valid_anchor or (requires_pr_fallback and not supports_suggestions_artifact): fallback_comments.append(f"{body}\n\nLocation: `{relevant_file}:" f"{relevant_lines_start}-{relevant_lines_end}`") else: - code_suggestions.append({'body': body, 'relevant_file': relevant_file, - 'relevant_lines_start': relevant_lines_start, - 'relevant_lines_end': relevant_lines_end, - 'original_suggestion': d}) + code_suggestions.append(rendered_suggestion) - if code_suggestions: + suggestions_to_publish = artifact_suggestions if supports_suggestions_artifact else code_suggestions + if suggestions_to_publish: if supports_suggestions_artifact: is_successful = self.git_provider.publish_code_suggestions_artifact( - code_suggestions, artifact_footer=coverage_footer) + suggestions_to_publish, artifact_footer=coverage_footer) + artifact_batch_published = is_successful else: is_successful = self.git_provider.publish_code_suggestions(code_suggestions) if is_successful: @@ -1152,7 +1278,7 @@ async def push_inline_code_suggestions(self, data, include_coverage_footer: bool self._output_published = True if coverage_footer and not supports_suggestions_artifact: fallback_comments.append(coverage_footer.strip()) - if fallback_comments: + if fallback_comments and not artifact_batch_published: self.git_provider.publish_comment("\n\n---\n\n".join(fallback_comments)) self._output_published = True if code_suggestions and not is_successful: @@ -1228,11 +1354,10 @@ def _get_patch_range_lines(patch, relevant_lines_start, relevant_lines_end) -> O target_line += 1 target_remaining -= 1 - if all(line_number in target_lines - for line_number in range(relevant_lines_start, relevant_lines_end + 1)): - return [target_lines[line_number] - for line_number in range(relevant_lines_start, relevant_lines_end + 1)] - return None + if len(target_lines) != relevant_lines_end - relevant_lines_start + 1: + return None + return [target_lines[line_number] + for line_number in range(relevant_lines_start, relevant_lines_end + 1)] def _validate_suggestion(self, relevant_file, relevant_lines_start, relevant_lines_end, existing_code) -> tuple[bool, str, bool]: @@ -1507,6 +1632,133 @@ def remove_line_numbers(self, patches_diff_list: List[str]) -> List[str]: get_logger().error(f"Error removing line numbers from patches_diff_list, error: {e}") return patches_diff_list + async def _predict_chunks(self, model: str, chunk_pairs: list) -> list: + if get_settings().pr_code_suggestions.parallel_calls: + results = await asyncio.gather( + *[self._get_prediction(model, numbered, unnumbered) for numbered, unnumbered in chunk_pairs], + return_exceptions=True, + ) + else: + results = [] + for numbered, unnumbered in chunk_pairs: + try: + results.append(await self._get_prediction(model, numbered, unnumbered)) + except Exception as error: + results.append(error) + for result in results: + if isinstance(result, BaseException) and not isinstance(result, Exception): + raise result + return results + + def _recovery_chain(self, model: str, settings) -> Optional[tuple]: + """Return the invocation-local fallback chain and current position after routing. + + The outer retry_with_fallback_models wrapper replaces the configured primary + pair with the routed one and records it as primary, so recovery must reproduce + that substitution; otherwise a routed run would search the configured chain + for a pair that is not there. Returns None when the caller's pair cannot be + found uniquely or a later pair repeats, in which case recovery gives up rather + than guess the position or retry an identical fallback. + """ + effective_chain = get_effective_fallback_chain() + if effective_chain is None: + get_logger().warning("Skipping chunk recovery: no active fallback invocation chain") + return None + original_deployment = settings.get("openai.deployment_id", None) + positions = [index for index, pair in enumerate(effective_chain) + if pair == (model, original_deployment)] + if len(positions) != 1: + get_logger().warning("Skipping chunk recovery: current model/deployment is not unique in the fallback chain") + return None + if len(set(effective_chain)) != len(effective_chain): + get_logger().warning("Skipping chunk recovery: the fallback chain repeats a model/deployment pair") + return None + return effective_chain, positions[0] + 1 + + async def _recover_failed_chunks(self, model: str, chunk_pairs: list, results: list) -> None: + """Try remaining models for failed slots, without replacing successful predictions.""" + # An entirely failed batch still belongs to the existing outer fallback loop. + if not any(isinstance(result, Exception) for result in results) or all( + isinstance(result, Exception) for result in results + ): + return + + settings = get_settings() + chain = self._recovery_chain(model, settings) + if chain is None: + return + fallback_chain, start = chain + original_deployment = settings.get("openai.deployment_id", None) + try: + for fallback_model, deployment in fallback_chain[start:]: + pending = [index for index, result in enumerate(results) if isinstance(result, Exception)] + if not pending: + break + # Keep the original diff intact; truncation must not imply complete coverage. + eligible = [] + recovered_any = False + # Resolve token controls under the same deployment that the fallback request will use. + settings.set("openai.deployment_id", deployment) + try: + token_handler = TokenHandler(model=fallback_model) + output_reserve = self._recovery_output_reserve(fallback_model) + budget = get_max_tokens(fallback_model) - output_reserve + for index in pending: + system, user = self._render_prediction_prompts(*chunk_pairs[index]) + tokens = token_handler.count_tokens(system) + token_handler.count_tokens(user) + if tokens <= budget: + eligible.append(index) + else: + get_logger().warning(f"Skipping recovery of chunk {index + 1} with {fallback_model}: " + "the complete prompt exceeds its token budget") + except Exception as error: + get_logger().warning(f"Cannot prepare chunk recovery with {fallback_model}: {error}") + continue + if not eligible: + continue + # Sibling calls have finished before the deployment switch above. + recovered = await self._predict_chunks(fallback_model, [chunk_pairs[index] for index in eligible]) + for index, result in zip(eligible, recovered, strict=True): + if isinstance(result, Exception): + get_logger().warning( + f"Failed to recover suggestion chunk {index + 1} with {fallback_model}", + artifact={"error": result}, + ) + continue + recovered_any = True + get_logger().info(f"Recovered suggestion chunk {index + 1} with {fallback_model}", + artifact={"error": results[index]}) + results[index] = result + if recovered_any: + # run_details' model line reports the outer primary unless a fallback + # ran, so mark this fallback usage stickily before it is overwritten. + record_model_used(fallback_model, is_fallback=True) + finally: + settings.set("openai.deployment_id", original_deployment) + + def _recovery_output_reserve(self, model: str) -> int: + """Return the completion headroom to reserve for a recovery fallback model. + + The LiteLLM handler reports the output allowance it will actually request + (config.max_output_tokens, extended thinking, per-provider controls); fall + back to the fixed soft threshold only when it exposes no specific limit so a + near-limit prompt is not classified as eligible without enough headroom. + """ + get_output_token_reserve = getattr(self.ai_handler, "get_output_token_reserve", None) + if callable(get_output_token_reserve): + try: + output_tokens = get_output_token_reserve(model, OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD) + except Exception as error: + get_logger().debug(f"Failed to resolve the output token reserve for {model}: {error}") + else: + if ( + isinstance(output_tokens, int) + and not isinstance(output_tokens, bool) + and output_tokens > 0 + ): + return output_tokens + return OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD + async def prepare_prediction_main(self, model: str) -> dict: self.failed_chunk_count = 0 self.total_chunk_count = 0 @@ -1548,38 +1800,17 @@ async def prepare_prediction_main(self, model: str) -> dict: zip(self.patches_diff_list, self.patches_diff_list_no_line_numbers, strict=True)) self.total_chunk_count = len(chunk_pairs) - # parallelize calls to AI: - if get_settings().pr_code_suggestions.parallel_calls: - prediction_results = await asyncio.gather( - *[self._get_prediction(model, patches_diff, patches_diff_no_line_numbers) for - patches_diff, patches_diff_no_line_numbers in chunk_pairs], - return_exceptions=True) - for chunk_index, prediction in enumerate(prediction_results): - if isinstance(prediction, Exception): - chunk_errors.append(prediction) - get_logger().warning( - f"Failed to generate code suggestions for chunk {chunk_index + 1}; " - "retaining successful chunks", - artifact={"error": prediction}, - ) - elif isinstance(prediction, BaseException): - raise prediction - else: - prediction_list.append(prediction) - else: - for chunk_index, (patches_diff, patches_diff_no_line_numbers) in enumerate( - chunk_pairs): - try: - prediction = await self._get_prediction(model, patches_diff, patches_diff_no_line_numbers) - except Exception as e: - chunk_errors.append(e) - get_logger().warning( - f"Failed to generate code suggestions for chunk {chunk_index + 1}; " - "retaining successful chunks", - artifact={"error": e}, - ) - else: - prediction_list.append(prediction) + prediction_results = await self._predict_chunks(model, chunk_pairs) + await self._recover_failed_chunks(model, chunk_pairs, prediction_results) + for chunk_index, prediction in enumerate(prediction_results): + if isinstance(prediction, Exception): + chunk_errors.append(prediction) + get_logger().warning( + f"Failed to generate code suggestions for chunk {chunk_index + 1}; retaining successful chunks", + artifact={"error": prediction}, + ) + else: + prediction_list.append(prediction) self.failed_chunk_count = len(chunk_errors) + self.parse_failure_count if chunk_errors and not prediction_list: @@ -1628,11 +1859,8 @@ async def convert_to_decoupled_with_line_numbers(self, patches_diff_list_no_line file=None).strip() patches_new[i] = patches_new[i].strip() patch_final = "\n\n\n".join(patches_new) - if model in MAX_TOKENS: - max_tokens_full = MAX_TOKENS[ - model] # note - here we take the actual max tokens, without any reductions. we do aim to get the full documentation website in the prompt - else: - max_tokens_full = get_max_tokens(model) + # take the actual max tokens, without any reductions + max_tokens_full = get_max_tokens(model, ignore_max_model_tokens=True) delta_output = 2000 token_count = self.token_handler.count_tokens(patch_final) if token_count > max_tokens_full - delta_output: @@ -1672,11 +1900,21 @@ def generate_summarized_suggestions(self, data: Dict) -> str: suggestions_labels = dict() # add all suggestions related to each label for suggestion in data['code_suggestions']: + if not self._is_suggestion_line_range_valid(suggestion): + # suggestions without resolved line anchors (e.g. when self-reflection + # failed or returned a mismatched count) cannot be placed in the diff; + # skip them instead of failing the whole table + continue label = suggestion['label'].strip().strip("'").strip('"') if label not in suggestions_labels: suggestions_labels[label] = [] suggestions_labels[label].append(suggestion) + if not suggestions_labels: + pr_body = f"{format_pr_code_suggestions_header()}\n\n" + pr_body += "No suggestions found to improve this PR." + return pr_body + # sort suggestions_labels by the suggestion with the highest score suggestions_labels = dict( sorted(suggestions_labels.items(), key=lambda x: max([s['score'] for s in x[1]]), reverse=True)) diff --git a/pr_agent/tools/pr_config.py b/pr_agent/tools/pr_config.py index 293db20e95..a4080f4643 100644 --- a/pr_agent/tools/pr_config.py +++ b/pr_agent/tools/pr_config.py @@ -15,7 +15,8 @@ def __init__(self, pr_url: str, args=None, ai_handler=None): Args: pr_url (str): The URL of the pull request to be reviewed. - args (list, optional): List of arguments passed to the PRReviewer class. Defaults to None. + args (list, optional): Unused compatibility parameter for the common tool interface. Defaults to None. + ai_handler: Unused compatibility parameter for the common tool interface. Defaults to None. """ self.git_provider = get_git_provider()(pr_url) diff --git a/pr_agent/tools/pr_description.py b/pr_agent/tools/pr_description.py index b1dc807a9e..971fb5de80 100644 --- a/pr_agent/tools/pr_description.py +++ b/pr_agent/tools/pr_description.py @@ -34,7 +34,7 @@ show_run_details, ) from pr_agent.config_loader import get_settings -from pr_agent.git_providers import GithubProvider, get_git_provider_with_context +from pr_agent.git_providers import get_git_provider_with_context from pr_agent.git_providers.git_provider import get_main_pr_language from pr_agent.log import get_logger from pr_agent.servers.help import HelpMessage @@ -43,6 +43,8 @@ fit_related_tickets_to_prompt_budget, ) +MAX_DESCRIPTION_COVERAGE_FILES = 50 + class PRDescription: def __init__(self, pr_url: str, args: list = None, @@ -136,19 +138,18 @@ async def run(self): if get_settings().pr_description.enable_semantic_files_types: self.file_label_dict = self._prepare_file_labels() - pr_labels, pr_file_changes = [], [] + pr_labels = [] if get_settings().pr_description.publish_labels: pr_labels = self._prepare_labels() else: get_logger().debug("Publishing labels disabled") if get_settings().pr_description.use_description_markers: - pr_title, pr_body, changes_walkthrough, pr_file_changes = self._prepare_pr_answer_with_markers() + pr_title, pr_body = self._prepare_pr_answer_with_markers() else: - pr_title, pr_body, changes_walkthrough, pr_file_changes = self._prepare_pr_answer() - if not self.git_provider.is_supported( - "publish_file_comments") or not get_settings().pr_description.inline_file_summary: - pr_body += "\n\n" + changes_walkthrough + "___\n\n" + pr_title, pr_body, changes_walkthrough = self._prepare_pr_answer() + pr_body += "\n\n" + changes_walkthrough + "___\n\n" + pr_body += self._get_description_coverage_footer() get_logger().debug("PR output", artifact={"title": pr_title, "body": pr_body}) # Add help text if gfm_markdown is supported @@ -157,15 +158,15 @@ async def run(self): pr_body += HelpMessage.get_describe_usage_guide() pr_body += "\n
      \n" elif get_settings().pr_description.enable_help_comment and self.git_provider.is_supported("gfm_markdown"): - if isinstance(self.git_provider, GithubProvider): + if self.git_provider.supports_inline_help_footer(): pr_body += ('\n\n___\n\n>
      Need help?
    • Type /help how to ... ' 'in the comments thread for any questions about PR-Agent usage.
    • Check out the ' - 'documentation ' + 'documentation ' 'for more information.
      • ') - else: # gitlab + else: # bullets separated by
      , for providers whose footer cannot inline a list pr_body += ("\n\n___\n\n
      Need help?- Type /help how to ... in the comments " "thread for any questions about PR-Agent usage.
      - Check out the " - "documentation for more information.
      ") + "documentation for more information.
      ") # elif get_settings().pr_description.enable_help_comment: # pr_body += '\n\n___\n\n> 💡 **PR-Agent usage**: Comment `/help "your question"` on any pull request to receive relevant information' @@ -259,6 +260,9 @@ async def run(self): return "" async def _prepare_prediction(self, model: str) -> None: + self.description_total_chunk_count = 0 + self.description_failed_chunk_count = 0 + self.description_failed_files = [] if get_settings().pr_description.use_description_markers and 'pr_agent:' not in self.user_description: get_logger().info("Markers were enabled, but user description does not contain markers. Skipping AI prediction") return None @@ -309,33 +313,87 @@ async def _prepare_prediction(self, model: str) -> None: self.git_provider, token_handler_only_files_prompt, model) # get the files prediction for each patch + chunk_pairs = list(zip(patches_compressed_list, files_in_patches_list, strict=True)) + self.description_total_chunk_count = len(chunk_pairs) + results = [None] * len(chunk_pairs) if not get_settings().pr_description.get("async_ai_calls", True): - results = [] - for i, patches in enumerate(patches_compressed_list): # sync calls + for i, (patches, _files_in_patch) in enumerate(chunk_pairs): # sync calls + if not patches: + continue patches_diff = "\n".join(patches) get_logger().debug(f"PR diff number {i + 1} for describe files") - prediction_files = await self._get_prediction(model, patches_diff, - prompt="pr_description_only_files_prompts") - results.append(prediction_files) + try: + results[i] = await self._get_prediction( + model, patches_diff, prompt="pr_description_only_files_prompts") + except Exception as e: + results[i] = e else: # async calls tasks = [] - for i, patches in enumerate(patches_compressed_list): + task_indices = [] + for i, (patches, _files_in_patch) in enumerate(chunk_pairs): if patches: patches_diff = "\n".join(patches) get_logger().debug(f"PR diff number {i + 1} for describe files") task = asyncio.create_task( self._get_prediction(model, patches_diff, prompt="pr_description_only_files_prompts")) tasks.append(task) + task_indices.append(i) # Wait for all tasks to complete - results = await asyncio.gather(*tasks) + task_results = await asyncio.gather(*tasks, return_exceptions=True) + for chunk_index, result in zip(task_indices, task_results, strict=True): + results[chunk_index] = result file_description_str_list = [] - for i, result in enumerate(results): + chunk_errors = [] + failed_files = [] + for i, (result, (_patches, files_in_patch)) in enumerate(zip(results, chunk_pairs, strict=True)): + if isinstance(result, Exception): + chunk_errors.append(result) + failed_files.extend(files_in_patch) + get_logger().warning( + f"Failed to generate description for chunk {i + 1}; retaining successful chunks", + artifact={"error": result, "files": files_in_patch}, + ) + continue + if isinstance(result, BaseException): + raise result + if not isinstance(result, str): + chunk_errors.append(ValueError(f"Description chunk {i + 1} returned no prediction")) + failed_files.extend(files_in_patch) + get_logger().warning( + f"Description chunk {i + 1} returned no prediction; retaining successful chunks", + artifact={"files": files_in_patch}, + ) + continue prediction_files = result.strip().removeprefix('```yaml').strip('`').strip() - if load_yaml(prediction_files, keys_fix_yaml=self.keys_fix) and prediction_files.startswith('pr_files'): + prediction_files_data = load_yaml(prediction_files, keys_fix_yaml=self.keys_fix) + file_descriptions = (prediction_files_data.get('pr_files') + if isinstance(prediction_files_data, dict) else None) + required_fields = ['filename', 'changes_title', 'label'] + if self.vars.get('include_file_summary_changes', True): + required_fields.append('changes_summary') + valid_file_descriptions = ( + isinstance(file_descriptions, list) and file_descriptions and + any(isinstance(file_description, dict) and + all(isinstance(file_description.get(field), str) and file_description[field].strip() + for field in required_fields) + for file_description in file_descriptions) + ) + if (prediction_files.startswith('pr_files') and isinstance(prediction_files_data, dict) and + valid_file_descriptions): prediction_files = prediction_files.removeprefix('pr_files:').strip() file_description_str_list.append(prediction_files) else: - get_logger().debug(f"failed to generate predictions in iteration {i + 1} for describe files") + chunk_errors.append(ValueError(f"Description chunk {i + 1} returned invalid YAML")) + failed_files.extend(files_in_patch) + get_logger().warning( + f"Failed to parse description chunk {i + 1}; retaining successful chunks", + artifact={"files": files_in_patch}, + ) + + self.description_failed_chunk_count = len(chunk_pairs) - len(file_description_str_list) + self.description_failed_files = list(dict.fromkeys(failed_files)) + if not file_description_str_list: + raise chunk_errors[0] if chunk_errors else ValueError("No description chunks were generated") # generate files_walkthrough string, with proper token handling self.vars, token_handler_only_description_prompt = fit_related_tickets_to_prompt_budget( @@ -390,6 +448,27 @@ async def _prepare_prediction(self, model: str) -> None: get_logger().debug(f"Using only headers for describe {self.pr_id}") self.prediction = prediction_headers + def _get_description_coverage_footer(self) -> str: + failed_chunk_count = getattr(self, "description_failed_chunk_count", 0) + if not failed_chunk_count: + return "" + + total_chunk_count = getattr(self, "description_total_chunk_count", failed_chunk_count) + footer = ( + f"\n\n⚠️ **Description coverage:** {failed_chunk_count} of {total_chunk_count} " + "file-description chunks failed; the description above is based on the successful chunks only." + ) + failed_files = getattr(self, "description_failed_files", []) + if failed_files: + displayed_files = failed_files[:MAX_DESCRIPTION_COVERAGE_FILES] + footer += "\n\nFiles from failed chunks may not be covered:\n" + "\n".join( + f"- `{filename}`" for filename in displayed_files + ) + remaining_count = len(failed_files) - len(displayed_files) + if remaining_count: + footer += f"\n... and {remaining_count} more" + return footer + async def extend_uncovered_files(self, original_prediction: str) -> str: try: prediction = original_prediction @@ -462,35 +541,6 @@ async def extend_uncovered_files(self, original_prediction: str) -> str: return original_prediction - async def extend_additional_files(self, remaining_files_list) -> str: - prediction = self.prediction - try: - original_prediction_dict = load_yaml(self.prediction, keys_fix_yaml=self.keys_fix) - prediction_extra = "pr_files:" - for file in remaining_files_list: - extra_file_yaml = f"""\ -- filename: | - {file} - changes_summary: | - ... - changes_title: | - ... - label: | - additional files (token-limit) -""" - prediction_extra = prediction_extra + "\n" + extra_file_yaml.strip() - prediction_extra_dict = load_yaml(prediction_extra, keys_fix_yaml=self.keys_fix) - # merge the two dictionaries - if isinstance(original_prediction_dict, dict) and isinstance(prediction_extra_dict, dict): - original_prediction_dict["pr_files"].extend(prediction_extra_dict["pr_files"]) - new_yaml = yaml.dump(original_prediction_dict) - if load_yaml(new_yaml, keys_fix_yaml=self.keys_fix): - prediction = new_yaml - return prediction - except Exception as e: - get_logger().error(f"Error extending additional files {self.pr_id}: {e}") - return self.prediction - async def _get_prediction(self, model: str, patches_diff: str, prompt="pr_description_prompt") -> str: variables = copy.deepcopy(self.vars) variables["diff"] = patches_diff # update diff @@ -568,7 +618,7 @@ def _prepare_labels(self) -> List[str]: get_logger().error(f"Error converting labels to original case {self.pr_id}: {e}") return pr_labels - def _prepare_pr_answer_with_markers(self) -> Tuple[str, str, str, List[dict]]: + def _prepare_pr_answer_with_markers(self) -> Tuple[str, str]: get_logger().info(f"Using description marker replacements {self.pr_id}") # Remove the 'PR Title' key from the dictionary @@ -609,11 +659,9 @@ def _prepare_pr_answer_with_markers(self) -> Tuple[str, str, str, List[dict]]: ai_walkthrough = self.data.get('pr_files') walkthrough_gfm = "" - pr_file_changes = [] if ai_walkthrough and not re.search(r'', body): try: - walkthrough_gfm, pr_file_changes = self.process_pr_files_prediction(walkthrough_gfm, - self.file_label_dict) + walkthrough_gfm = self.process_pr_files_prediction(walkthrough_gfm, self.file_label_dict) body = body.replace('pr_agent:walkthrough', walkthrough_gfm) except Exception as e: get_logger().error(f"Failing to process walkthrough {self.pr_id}: {e}") @@ -624,9 +672,9 @@ def _prepare_pr_answer_with_markers(self) -> Tuple[str, str, str, List[dict]]: if ai_diagram: body = re.sub(r'|pr_agent:diagram', ai_diagram, body) - return title, body, walkthrough_gfm, pr_file_changes + return title, body - def _prepare_pr_answer(self) -> Tuple[str, str, str, List[dict]]: + def _prepare_pr_answer(self) -> Tuple[str, str, str]: """ Prepare the PR description based on the AI prediction data. @@ -656,7 +704,6 @@ def _prepare_pr_answer(self) -> Tuple[str, str, str, List[dict]]: # Iterate over the remaining dictionary items and append the key and value to 'pr_body' in a markdown format, # except for the items containing the word 'walkthrough' pr_body, changes_walkthrough = "", "" - pr_file_changes = [] for idx, (key, value) in enumerate(self.data.items()): if key == 'changes_diagram': pr_body += f"### {PRDescriptionHeader.DIAGRAM_WALKTHROUGH.value}\n\n" @@ -681,7 +728,7 @@ def _prepare_pr_answer(self) -> Tuple[str, str, str, List[dict]]: if self.git_provider.is_supported("gfm_markdown"): pr_body += "
      • \n" elif 'pr_files' in key.lower() and get_settings().pr_description.enable_semantic_files_types: # 'File Walkthrough' section - changes_walkthrough_table, pr_file_changes = self.process_pr_files_prediction(changes_walkthrough, value) + changes_walkthrough_table = self.process_pr_files_prediction(changes_walkthrough, value) if get_settings().pr_description.get('file_table_collapsible_open_by_default', False): initial_status = " open" else: @@ -702,7 +749,7 @@ def _prepare_pr_answer(self) -> Tuple[str, str, str, List[dict]]: if idx < len(self.data) - 1: pr_body += "\n\n___\n\n" - return title, pr_body, changes_walkthrough, pr_file_changes, + return title, pr_body, changes_walkthrough def _prepare_file_labels(self): file_label_dict = {} @@ -739,7 +786,6 @@ def _prepare_file_labels(self): return file_label_dict def process_pr_files_prediction(self, pr_body, value): - pr_comments = [] # logic for using collapsible file list use_collapsible_file_list = get_settings().pr_description.collapsible_file_list num_files = 0 @@ -750,7 +796,7 @@ def process_pr_files_prediction(self, pr_body, value): use_collapsible_file_list = num_files > self.COLLAPSIBLE_FILE_LIST_THRESHOLD if not self.git_provider.is_supported("gfm_markdown"): - return pr_body, pr_comments + return pr_body try: pr_body += "" header = "Relevant files" @@ -816,7 +862,7 @@ def process_pr_files_prediction(self, pr_body, value): except Exception as e: get_logger().error(f"Error processing pr files to markdown {self.pr_id}: {str(e)}") pass - return pr_body, pr_comments + return pr_body def add_file_data(self, delta_nbsp, diff_plus_minus, file_change_description_br, filename, filename_publish, link, pr_body) -> str: diff --git a/pr_agent/tools/pr_help_docs.py b/pr_agent/tools/pr_help_docs.py index ffcdb72cc0..5f7afbba14 100644 --- a/pr_agent/tools/pr_help_docs.py +++ b/pr_agent/tools/pr_help_docs.py @@ -7,7 +7,6 @@ from jinja2 import Environment, StrictUndefined -from pr_agent.algo import MAX_TOKENS from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler from pr_agent.algo.pr_processing import retry_with_fallback_models @@ -69,6 +68,9 @@ def extract_model_answer_and_relevant_sources(ai_response: str) -> str | None: get_logger().info(f"Found model answer: {model_answer_section_in_response}") return model_answer_and_relevant_sources_sections_in_response \ if len(model_answer_section_in_response) > 0 else None + if model_answer_and_relevant_sources_sections_in_response.strip(): + get_logger().info(f"Found model answer without relevant sources: {model_answer_and_relevant_sources_sections_in_response}") + return model_answer_and_relevant_sources_sections_in_response get_logger().warning(f"Either no answer section found, or that section is malformed: {ai_response}") return None @@ -169,23 +171,39 @@ def aggregate_documentation_files_for_prompt_contents(file_path_to_contents: dic def format_markdown_q_and_a_response(question_str: str, response_str: str, relevant_sections: list[dict[str, str]], supported_suffixes: list[str], base_url_prefix: str, base_url_suffix: str="") -> str: try: - base_url_prefix = base_url_prefix.strip('/') #Sanitize base_url_prefix + base_url_prefix = base_url_prefix.strip('/') # Sanitize base_url_prefix answer_str = "" answer_str += f"### Question: \n{question_str}\n\n" answer_str += f"### Answer:\n{response_str.strip()}\n\n" - answer_str += "#### Relevant Sources:\n\n" + source_links = [] + if not isinstance(relevant_sections, list): + get_logger().warning("Skipping malformed relevant source collection: expected a list") + relevant_sections = [] for section in relevant_sections: - file = section.get('file_name').lstrip('/').strip() #Remove any '/' in the beginning, since some models do it anyway + try: + if not isinstance(section, dict): + raise ValueError("source row is not a mapping") + file_name = section.get('file_name') + header = section.get('relevant_section_header_string') + if not isinstance(file_name, str) or not isinstance(header, str): + raise ValueError("source row has an invalid file name or section header") + except ValueError as e: + get_logger().warning(f"Skipping malformed relevant source row: {e}") + continue + file = file_name.lstrip('/').strip() # Remove any leading '/', since some models add one ext = [suffix for suffix in supported_suffixes if file.endswith(suffix)] if not ext: get_logger().warning(f"Unsupported file extension: {file}") continue - if str(section['relevant_section_header_string']).strip(): - markdown_header = format_markdown_header(section['relevant_section_header_string']) + if header.strip(): + markdown_header = format_markdown_header(header) if base_url_prefix: - answer_str += f"> - {base_url_prefix}/{file}{base_url_suffix}#{markdown_header}\n" + source_links.append(f"> - {base_url_prefix}/{file}{base_url_suffix}#{markdown_header}\n") else: - answer_str += f"> - {base_url_prefix}/{file}{base_url_suffix}\n" + source_links.append(f"> - {base_url_prefix}/{file}{base_url_suffix}\n") + if source_links: + answer_str += "#### Relevant Sources:\n\n" + answer_str += ''.join(source_links) return answer_str except Exception as e: get_logger().exception("Unexpected exception thrown. Returning empty result.") @@ -218,6 +236,29 @@ def format_markdown_header(header: str) -> str: get_logger().exception("Error while formatting markdown header", artifacts={'header': header}) return "" + +def get_valid_ranking_indices(ranking: list[dict], number_of_documents: int) -> list[int]: + """Return in-range document indices while isolating malformed ranking rows.""" + valid_indices = [] + if not isinstance(ranking, list): + get_logger().warning("Skipping malformed relevant file ranking collection: expected a list") + return valid_indices + for entry in ranking: + try: + if not isinstance(entry, dict): + raise ValueError("ranking row is not a mapping") + raw_index = entry.get('idx') + if isinstance(raw_index, bool) or not isinstance(raw_index, (int, str)): + raise ValueError("ranking index has an invalid type") + index = int(raw_index) + if 0 <= index < number_of_documents: + valid_indices.append(index) + else: + raise ValueError("ranking index is out of range") + except (TypeError, ValueError) as e: + get_logger().warning(f"Skipping malformed relevant file ranking row: {e}") + return valid_indices + def clean_markdown_content(content: str) -> str: """ Remove hidden comments and unnecessary elements from markdown content to reduce size. @@ -371,9 +412,11 @@ async def run(self): return response_str = response_yaml.get('response') relevant_sections = response_yaml.get('relevant_sections') - if not response_str or not relevant_sections: - get_logger().error("Failed to extract response/relevant sections.", - artifacts={'raw_response': response, 'response_str': response_str, 'relevant_sections': relevant_sections}) + if not response_str: + get_logger().error( + "Failed to extract response.", + artifacts={'raw_response': response, 'response_str': response_str, + 'relevant_sections': relevant_sections}) return if int(response_yaml.get('question_is_relevant', '1')) == 0: get_logger().warning("Question is not relevant. Returning without an answer...", @@ -418,6 +461,24 @@ def _find_all_document_files_matching_exts(self, abs_docs_path: str, ignore_read get_logger().exception("Unexpected exception thrown. Returning empty list.") return [] + def _filter_doc_files_under_clone(self, doc_files: list[str], clone_root: str) -> list[str]: + """Drop candidates whose resolved path leaves the cloned repository root. + + Symlinked documentation files can point outside the clone at host files; resolving + each candidate keeps anything that escapes from ever being read into the model prompt. + """ + contained = [] + for file_path in doc_files: + resolved = os.path.realpath(file_path) + if resolved != clone_root and not resolved.startswith(clone_root + os.sep): + get_logger().warning( + f"Documentation file '{file_path}' resolves outside the cloned repository " + "root; skipping it" + ) + continue + contained.append(file_path) + return contained + def _gen_filenames_to_contents_map_from_repo(self) -> dict[str, str]: try: with TemporaryDirectory() as tmp_dir: @@ -435,15 +496,29 @@ def _gen_filenames_to_contents_map_from_repo(self) -> dict[str, str]: for file in files: if file.lower().startswith("readme."): doc_files.append(os.path.join(root, file)) - abs_docs_path = os.path.join(returned_cloned_repo_root.path, self.docs_path) - if os.path.exists(abs_docs_path): - doc_files.extend(self._find_all_document_files_matching_exts(abs_docs_path, - ignore_readme=(self.docs_path=='.'))) - if not doc_files: - get_logger().warning(f"No documentation files found matching file extensions: " - f"{self.supported_doc_exts} under repo: {self.repo_url} " - f"path: {self.docs_path}. Returning empty dict.") - return {} + # Resolve docs_path against the clone root and refuse paths that escape it: + # an absolute or traversing docs_path (e.g. an attacker-controlled per-directory + # setting) must never read host files into the prompt. + clone_root = os.path.realpath(returned_cloned_repo_root.path) + resolved_docs_path = os.path.realpath(os.path.join(clone_root, self.docs_path)) + if resolved_docs_path == clone_root or resolved_docs_path.startswith(clone_root + os.sep): + if os.path.exists(resolved_docs_path): + doc_files.extend(self._find_all_document_files_matching_exts(resolved_docs_path, + ignore_readme=(self.docs_path=='.'))) + else: + get_logger().warning( + f"docs_path '{self.docs_path}' escapes the cloned repository root; " + "skipping docs-path file gathering" + ) + # Individual files can still escape the clone through symlinks (a doc or root + # README that points at a host file). Resolve every candidate and refuse + # anything that leaves the canonical clone root so host files never reach the prompt. + doc_files = self._filter_doc_files_under_clone(doc_files, clone_root) + if not doc_files: + get_logger().warning(f"No documentation files found matching file extensions: " + f"{self.supported_doc_exts} under repo: {self.repo_url} " + f"path: {self.docs_path}. Returning empty dict.") + return {} get_logger().info(f'For context {self.ctx_url} and repo: {self.repo_url}' f' will be using the following documentation files: ', @@ -465,12 +540,9 @@ def _trim_docs_input(self, docs_input: str, max_allowed_txt_input: int, only_ret # Then, count the tokens in the prompt. If the count exceeds the limit, trim the text. token_count = self.token_handler.count_tokens(docs_input, force_accurate=True) get_logger().debug(f"Estimated token count of documentation to send to model: {token_count}") - model = get_settings().config.model - if model in MAX_TOKENS: - max_tokens_full = MAX_TOKENS[ - model] # note - here we take the actual max tokens, without any reductions. we do aim to get the full documentation website in the prompt - else: - max_tokens_full = get_max_tokens(model) + # take the actual max tokens, without any reductions. we do aim to get + # the full documentation website in the prompt + max_tokens_full = get_max_tokens(get_settings().config.model, ignore_max_model_tokens=True) delta_output = 5000 # Elbow room to reduce chance of exceeding token limit or model paying less attention to prompt guidelines. if token_count > max_tokens_full - delta_output: if only_return_if_trim_needed: @@ -515,8 +587,8 @@ async def _rank_docs_and_return_them_as_prompt(self, docs_filepath_to_contents: get_logger().error("Failed to parse the AI response.", artifacts={'response': response}) return "" # else: Sanitize the output so that the file names match 1:1 dictionary keys. Do this via the file index and not its name, which may be altered by the model. - valid_indices = [int(entry['idx']) for entry in response_yaml.get('relevant_files_ranking') - if int(entry['idx']) >= 0 and int(entry['idx']) < len(docs_filepath_to_contents)] + valid_indices = get_valid_ranking_indices(response_yaml.get('relevant_files_ranking'), + len(docs_filepath_to_contents)) valid_file_paths = [list(docs_filepath_to_contents.keys())[idx] for idx in valid_indices] selected_docs_dict = {file_path: docs_filepath_to_contents[file_path] for file_path in valid_file_paths} docs_prompt = aggregate_documentation_files_for_prompt_contents(selected_docs_dict) diff --git a/pr_agent/tools/pr_help_message.py b/pr_agent/tools/pr_help_message.py index ae767aebd4..70c1f7c881 100644 --- a/pr_agent/tools/pr_help_message.py +++ b/pr_agent/tools/pr_help_message.py @@ -1,34 +1,28 @@ import copy import re from functools import partial +from math import ceil, isfinite from pathlib import Path -from jinja2 import Environment, StrictUndefined +from jinja2 import Environment, StrictUndefined, select_autoescape +from litellm import token_counter -from pr_agent.algo import MAX_TOKENS from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler from pr_agent.algo.pr_processing import retry_with_fallback_models -from pr_agent.algo.token_handler import TokenHandler -from pr_agent.algo.utils import ModelType, clip_tokens, get_max_tokens, load_yaml +from pr_agent.algo.token_handler import TokenEncoder +from pr_agent.algo.utils import ModelType, get_max_tokens, load_yaml from pr_agent.command_descriptions import COMMAND_DESCRIPTIONS from pr_agent.config_loader import get_settings -from pr_agent.git_providers import BitbucketServerProvider, GithubProvider, get_git_provider_with_context +from pr_agent.git_providers import get_git_provider_with_context from pr_agent.log import get_logger +DOCS_SITE_URL = "https://docs.pr-agent.ai" +HELP_OUTPUT_TOKEN_RESERVE = 2_000 +MESSAGE_FRAMING_TOKEN_ALLOWANCE = 16 +REPLY_FRAMING_TOKEN_ALLOWANCE = 16 +TRUNCATION_MARKER = "\n...(truncated)\n" -def extract_header(snippet): - res = '' - lines = snippet.split('===Snippet content===')[0].split('\n') - highest_header = '' - highest_level = float('inf') - for line in lines[::-1]: - line = line.strip() - if line.startswith('Header '): - highest_header = line.split(': ')[1] - if highest_header: - res = f"#{highest_header.lower().replace(' ', '-')}" - return res class PRHelpMessage: def __init__(self, pr_url: str, args=None, ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler, return_as_string=False): @@ -41,23 +35,164 @@ def __init__(self, pr_url: str, args=None, ai_handler: partial[BaseAiHandler,] = "question": self.question_str, "snippets": "", } - self.token_handler = TokenHandler(None, - self.vars, - get_settings().pr_help_prompts.system, - get_settings().pr_help_prompts.user) async def _prepare_prediction(self, model: str): + variables = copy.deepcopy(self.vars) + system_prompt, user_prompt = self._fit_prompts(variables, model) + response, finish_reason = await self.ai_handler.chat_completion( + model=model, temperature=get_settings().config.temperature, system=system_prompt, user=user_prompt) + return response + + @staticmethod + def _render_prompts(variables): + # These string templates produce plain-text model prompts, not HTML. + environment = Environment( + autoescape=select_autoescape(default_for_string=False), + undefined=StrictUndefined, + ) + system_prompt = environment.from_string(get_settings().pr_help_prompts.system).render(variables) + user_prompt = environment.from_string(get_settings().pr_help_prompts.user).render(variables) + return system_prompt, user_prompt + + @staticmethod + def _coerce_output_token_limit(value) -> int: try: - variables = copy.deepcopy(self.vars) - environment = Environment(undefined=StrictUndefined) - system_prompt = environment.from_string(get_settings().pr_help_prompts.system).render(variables) - user_prompt = environment.from_string(get_settings().pr_help_prompts.user).render(variables) - response, finish_reason = await self.ai_handler.chat_completion( - model=model, temperature=get_settings().config.temperature, system=system_prompt, user=user_prompt) - return response + output_tokens = int(value) + except (TypeError, ValueError, OverflowError): + return 0 + return output_tokens if output_tokens > 0 else 0 + + def _get_prompt_budget(self, model: str) -> int: + output_tokens = 0 + get_output_token_reserve = getattr(self.ai_handler, "get_output_token_reserve", None) + if callable(get_output_token_reserve): + try: + handler_output_tokens = get_output_token_reserve(model, HELP_OUTPUT_TOKEN_RESERVE) + except Exception as e: + get_logger().debug(f"Failed to resolve the output token reserve for {model}: {e}") + else: + if ( + isinstance(handler_output_tokens, int) + and not isinstance(handler_output_tokens, bool) + and handler_output_tokens > 0 + ): + output_tokens = handler_output_tokens + get_output_token_limit = getattr(self.ai_handler, "get_output_token_limit", None) + if output_tokens <= 0 and callable(get_output_token_limit): + try: + handler_output_tokens = get_output_token_limit(model) + except Exception as e: + get_logger().debug(f"Failed to resolve the output token limit for {model}: {e}") + else: + if ( + isinstance(handler_output_tokens, int) + and not isinstance(handler_output_tokens, bool) + and handler_output_tokens > 0 + ): + output_tokens = handler_output_tokens + if output_tokens <= 0: + raw_output_tokens = get_settings().config.get("max_output_tokens", 0) + output_tokens = self._coerce_output_token_limit(raw_output_tokens) + if output_tokens <= 0: + output_tokens = HELP_OUTPUT_TOKEN_RESERVE + return max(get_max_tokens(model, ignore_max_model_tokens=True) - output_tokens, 0) + + @staticmethod + def _count_prompt_tokens(model: str, system_prompt: str, user_prompt: str) -> int: + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] + try: + model_token_count = token_counter(model=model, messages=messages) + if isinstance(model_token_count, int) and not isinstance(model_token_count, bool) and model_token_count > 0: + return model_token_count except Exception as e: - get_logger().error(f"Error while preparing prediction: {e}") - return "" + get_logger().debug(f"Model-aware token counting failed for {model}: {e}") + + get_logger().debug(f"Using a local token estimate for {model}") + encoder = TokenEncoder.get_token_encoder(model) + content_tokens = sum( + len(encoder.encode(message["content"], disallowed_special=())) for message in messages + ) + framing_tokens = MESSAGE_FRAMING_TOKEN_ALLOWANCE * len(messages) + REPLY_FRAMING_TOKEN_ALLOWANCE + raw_factor = get_settings().get("config.model_token_count_estimate_factor", 0) + try: + extra_factor = float(raw_factor) + except (TypeError, ValueError, OverflowError): + extra_factor = 0 + if isinstance(raw_factor, bool) or not isfinite(extra_factor): + extra_factor = 0 + multiplier = max(1.0, 1.0 + extra_factor) + raw_estimate = content_tokens + framing_tokens + try: + estimated_tokens = raw_estimate * multiplier + if not isfinite(estimated_tokens): + raise ValueError("non-finite token estimate") + return ceil(estimated_tokens) + except (OverflowError, ValueError): + get_logger().warning( + f"model_token_count_estimate_factor is too large ({raw_factor!r}), using the estimate as is" + ) + return raw_estimate + + def _normalize_request_prompts(self, model: str, system_prompt: str, user_prompt: str) -> tuple[str, str]: + normalize_request_prompts = getattr(self.ai_handler, "normalize_request_prompts", None) + if not callable(normalize_request_prompts): + return system_prompt, user_prompt + try: + normalized_prompts = normalize_request_prompts(model, system_prompt, user_prompt) + except Exception as e: + get_logger().debug(f"Failed to normalize prompts for {model}: {e}") + return system_prompt, user_prompt + if ( + isinstance(normalized_prompts, tuple) + and len(normalized_prompts) == 2 + and all(isinstance(prompt, str) for prompt in normalized_prompts) + ): + return normalized_prompts + get_logger().debug(f"Ignoring unusable prompt normalization result for {model}") + return system_prompt, user_prompt + + def _fit_prompts(self, variables, model: str): + prompt_budget = self._get_prompt_budget(model) + raw_snippets = variables.get("snippets", "") + + def render(snippets): + attempt_variables = copy.deepcopy(variables) + attempt_variables["snippets"] = snippets + rendered_prompts = self._render_prompts(attempt_variables) + return self._normalize_request_prompts(model, *rendered_prompts) + + full_prompts = render(raw_snippets) + if self._count_prompt_tokens(model, *full_prompts) <= prompt_budget: + return full_prompts + + empty_prompts = render("") + if self._count_prompt_tokens(model, *empty_prompts) > prompt_budget: + raise ValueError(f"The /help prompt exceeds the token limit for {model} without documentation") + + marker_prompts = render(TRUNCATION_MARKER) + if self._count_prompt_tokens(model, *marker_prompts) > prompt_budget: + raise ValueError(f"The /help prompt exceeds the token limit for {model} with a truncation marker") + + keep_chars = max(len(raw_snippets) - 1, 0) + while keep_chars > 0: + candidate = raw_snippets[:keep_chars].rstrip() + TRUNCATION_MARKER + candidate_prompts = render(candidate) + candidate_tokens = self._count_prompt_tokens(model, *candidate_prompts) + if candidate_tokens <= prompt_budget: + get_logger().warning( + f"Documentation was clipped for /help to fit the {prompt_budget}-token input limit for {model}" + ) + return candidate_prompts + next_keep_chars = max(1, int(keep_chars * prompt_budget / candidate_tokens)) + keep_chars = min(keep_chars - 1, next_keep_chars) + + get_logger().warning( + f"Documentation was clipped for /help to fit the {prompt_budget}-token input limit for {model}" + ) + return marker_prompts def parse_args(self, args): if args and len(args) > 0: @@ -93,26 +228,32 @@ def format_markdown_header(self, header: str) -> str: get_logger().exception("Error while formatting markdown header", artifacts={'header': header}) return "" + def format_docs_url(self, file_name: str, header: str) -> str: + relative_path = file_name.strip().lstrip('/').removesuffix('.md') + if relative_path == 'index': + relative_path = '' + elif relative_path.endswith('/index'): + relative_path = relative_path.removesuffix('index') + elif relative_path: + relative_path += '/' + + docs_url = f"{DOCS_SITE_URL}/{relative_path}" + if str(header).strip(): + docs_url += f"#{self.format_markdown_header(header)}" + return docs_url + async def run(self): try: if self.question_str: get_logger().info(f'Answering a PR question about the PR {self.git_provider.pr_url} ') - if not get_settings().get('openai.key'): - if get_settings().config.publish_output: - self.git_provider.publish_comment( - "The `Help` tool chat feature requires an OpenAI API key for calculating embeddings") - else: - get_logger().error("The `Help` tool chat feature requires an OpenAI API key for calculating embeddings") - return - # current path docs_path= Path(__file__).parent.parent.parent / 'docs' / 'docs' # get all the 'md' files inside docs_path and its subdirectories md_files = list(docs_path.glob('**/*.md')) folders_to_exclude = ['/finetuning_benchmark/'] - files_to_exclude = {'EXAMPLE_BEST_PRACTICE.md', 'compression_strategy.md', '/docs/overview/index.md'} + files_to_exclude = {'compression_strategy.md', '/docs/overview/index.md'} md_files = [file for file in md_files if not any(folder in str(file) for folder in folders_to_exclude) and not any(file.name == file_to_exclude for file_to_exclude in files_to_exclude)] # sort the 'md_files' so that 'priority_files' will be at the top @@ -131,18 +272,6 @@ async def run(self): docs_prompt += f"\n==file name==\n\n{file_path}\n\n==file content==\n\n{f.read().strip()}\n=========\n\n" except Exception as e: get_logger().error(f"Error while reading the file {file}: {e}") - token_count = self.token_handler.count_tokens(docs_prompt) - get_logger().debug(f"Token count of full documentation website: {token_count}") - - model = get_settings().config.model - if model in MAX_TOKENS: - max_tokens_full = MAX_TOKENS[model] # note - here we take the actual max tokens, without any reductions. we do aim to get the full documentation website in the prompt - else: - max_tokens_full = get_max_tokens(model) - delta_output = 2000 - if token_count > max_tokens_full - delta_output: - get_logger().info(f"Token count {token_count} exceeds the limit {max_tokens_full - delta_output}. Skipping the PR Help message.") - docs_prompt = clip_tokens(docs_prompt, max_tokens_full - delta_output) self.vars['snippets'] = docs_prompt.strip() # run the AI model @@ -174,14 +303,12 @@ async def run(self): answer_str += f"### Question: \n{self.question_str}\n\n" answer_str += f"### Answer:\n{response_str.strip()}\n\n" answer_str += "#### Relevant Sources:\n\n" - base_path = "https://qodo-merge-docs.qodo.ai/" for section in relevant_sections: - file = section.get('file_name').strip().removesuffix('.md') - if str(section['relevant_section_header_string']).strip(): - markdown_header = self.format_markdown_header(section['relevant_section_header_string']) - answer_str += f"> - {base_path}{file}#{markdown_header}\n" - else: - answer_str += f"> - {base_path}{file}\n" + docs_url = self.format_docs_url( + section.get('file_name'), + section['relevant_section_header_string'], + ) + answer_str += f"> - {docs_url}\n" # publish the answer @@ -190,19 +317,20 @@ async def run(self): else: get_logger().info(f"Answer:\n{answer_str}") else: - if not isinstance(self.git_provider, BitbucketServerProvider) and not self.git_provider.is_supported("gfm_markdown"): + supports_gfm_markdown = self.git_provider.is_supported("gfm_markdown") + if not supports_gfm_markdown and not self.git_provider.supports_markdown_tables(): self.git_provider.publish_comment( "The `Help` tool requires gfm markdown, which is not supported by your code platform.") return get_logger().info('Getting PR Help Message...') - relevant_configs = {'pr_help': dict(get_settings().pr_help), + relevant_configs = {'pr_help': dict(get_settings().get("pr_help", {})), 'config': dict(get_settings().config)} get_logger().debug("Relevant configs", artifacts=relevant_configs) pr_comment = "## PR Agent Walkthrough 🤖\n\n" pr_comment += "Welcome to the PR Agent, an AI-powered tool for automated pull request analysis, feedback, suggestions and more.""" pr_comment += "\n\nHere is a list of tools you can use to interact with the PR Agent:\n" - base_path = "https://pr-agent-docs.codium.ai/tools" + base_path = f"{DOCS_SITE_URL}/tools" tool_names = [] tool_names.append(f"[DESCRIBE]({base_path}/describe/)") @@ -242,22 +370,23 @@ async def run(self): checkbox_list.append("[*]") checkbox_list.append("[*]") - if isinstance(self.git_provider, GithubProvider) and not get_settings().config.get('disable_checkboxes', False): + if (supports_gfm_markdown and self.git_provider.supports_checkbox_commands() + and not get_settings().config.get('disable_checkboxes', False)): pr_comment += "
      " for i in range(len(tool_names)): pr_comment += f"\n\n\n" pr_comment += "
      ToolDescriptionTrigger Interactively :gem:
      \n\n{tool_names[i]}{descriptions[i]}\n\n{checkbox_list[i]}\n
      \n\n" - pr_comment += """\n\n(1) Note that each tool can be [triggered automatically](https://pr-agent-docs.codium.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) when a new PR is opened, or called manually by [commenting on a PR](https://pr-agent-docs.codium.ai/usage-guide/automations_and_usage/#online-usage).""" + pr_comment += """\n\n(1) Note that each tool can be [triggered automatically](https://docs.pr-agent.ai/usage-guide/automations_and_usage/#github-app-automatic-tools-when-a-new-pr-is-opened) when a new PR is opened, or called manually by [commenting on a PR](https://docs.pr-agent.ai/usage-guide/automations_and_usage/#online-usage).""" pr_comment += """\n\n(2) Tools marked with [*] require additional parameters to be passed. For example, to invoke the `/ask` tool, you need to comment on a PR: `/ask ""`. See the relevant documentation for each tool for more details.""" - elif isinstance(self.git_provider, BitbucketServerProvider): - # only support basic commands in BBDC + elif not supports_gfm_markdown: + # only basic commands, in a plain markdown table (e.g. BBDC) pr_comment = generate_bbdc_table(tool_names[:4], descriptions[:4]) else: pr_comment += "" for i in range(len(tool_names)): pr_comment += f"\n" pr_comment += "
      ToolCommandDescription
      \n\n{tool_names[i]}{commands[i]}{descriptions[i]}
      \n\n" - pr_comment += """\n\nNote that each tool can be [invoked automatically](https://pr-agent-docs.codium.ai/usage-guide/automations_and_usage/) when a new PR is opened, or called manually by [commenting on a PR](https://pr-agent-docs.codium.ai/usage-guide/automations_and_usage/#online-usage).""" + pr_comment += """\n\nNote that each tool can be [invoked automatically](https://docs.pr-agent.ai/usage-guide/automations_and_usage/) when a new PR is opened, or called manually by [commenting on a PR](https://docs.pr-agent.ai/usage-guide/automations_and_usage/#online-usage).""" if get_settings().config.publish_output: self.git_provider.publish_comment(pr_comment) @@ -265,26 +394,6 @@ async def run(self): get_logger().exception(f"Error while running PRHelpMessage: {e}") return "" - async def prepare_relevant_snippets(self, sim_results): - # Get relevant snippets - relevant_snippets_full = [] - relevant_pages_full = [] - relevant_snippets_full_header = [] - th = 0.75 - for s in sim_results: - page = s[0].metadata['source'] - content = s[0].page_content - score = s[1] - relevant_snippets_full.append(content) - relevant_snippets_full_header.append(extract_header(content)) - relevant_pages_full.append(page) - # build the snippets string - relevant_snippets_str = "" - for i, s in enumerate(relevant_snippets_full): - relevant_snippets_str += f"Snippet {i+1}:\n\n{s}\n\n" - relevant_snippets_str += "-------------------\n\n" - return relevant_pages_full, relevant_snippets_full_header, relevant_snippets_str - def generate_bbdc_table(column_arr_1, column_arr_2): # Generating header row diff --git a/pr_agent/tools/pr_reviewer.py b/pr_agent/tools/pr_reviewer.py index 529230b862..10d9155ebf 100644 --- a/pr_agent/tools/pr_reviewer.py +++ b/pr_agent/tools/pr_reviewer.py @@ -18,6 +18,7 @@ key_issue_location_fingerprint, ) from pr_agent.algo.pr_processing import ( + PreparedPRDiff, add_ai_metadata_to_diff_files, get_pr_diff, get_pr_multi_diffs, @@ -462,10 +463,6 @@ def _review_finding_state_enabled(self) -> bool: provider = getattr(self, "git_provider", None) if provider is None: return False - publisher = getattr(provider, "publish_persistent_comment", None) - if getattr(publisher, "__func__", None) is GitProvider.publish_persistent_comment: - # Skip generic publishers; they only create comments and cannot safely carry lifecycle state. - return False if ( getattr(getattr(settings, "github", None), "publish_as_check_run", False) and callable(getattr(provider, "_publish_check_run", None)) @@ -731,6 +728,9 @@ def _prepare_review_finding_state(self, data: dict) -> None: allow_resolution = ( bool(self.prediction) and not bool(getattr(self.incremental, "is_incremental", False)) + # A merged result with failed chunks is still partial, even when chunking left + # no additional token-budget files to report. + and not bool(self.review_failed_chunk_count) and not bool(self.remaining_files_list) and parsed.valid and current_findings is not None @@ -753,6 +753,12 @@ def _should_publish_review_no_suggestions(self, pr_review: str) -> bool: return get_settings().pr_reviewer.get('publish_output_no_suggestions', True) or "No major issues detected" not in pr_review async def _prepare_prediction(self, model: str) -> None: + # Each model attempt owns a fresh result. A malformed primary must not + # leave state that can be mistaken for a successful fallback response. + self.prediction = None + self.prediction_data = None + self.review_chunk_count = 1 + self.review_failed_chunk_count = 0 raw_prompt_vars = getattr(self, "_raw_prompt_vars", getattr(self, "vars", None)) if raw_prompt_vars is not None: self.vars, self.token_handler = fit_related_tickets_to_prompt_budget( @@ -762,54 +768,76 @@ async def _prepare_prediction(self, model: str) -> None: get_settings().pr_review_prompt.user, model, ) - output = get_pr_diff(self.git_provider, - self.token_handler, - model, - add_line_numbers_to_hunks=True, - disable_extra_lines=False, - return_remaining_files=True,) - if isinstance(output, tuple): + chunking_enabled = get_settings().pr_reviewer.get("enable_large_pr_chunking", False) + diff_kwargs = { + "add_line_numbers_to_hunks": True, + "disable_extra_lines": False, + "return_remaining_files": True, + } + if chunking_enabled: + diff_kwargs["return_prepared"] = True + output = get_pr_diff(self.git_provider, self.token_handler, model, **diff_kwargs) + if isinstance(output, PreparedPRDiff): + self.patches_diff = output.diff + self.remaining_files_list = output.remaining_files_list + elif isinstance(output, tuple): self.patches_diff, self.remaining_files_list = output else: self.patches_diff = output self.remaining_files_list = [] # a non-empty remaining_files_list means the token budget truncated the diff - if self.remaining_files_list and get_settings().pr_reviewer.get("enable_large_pr_chunking", False): - if await self._prepare_chunked_prediction(model): + if self.remaining_files_list and chunking_enabled: + prepared_diff = output if isinstance(output, PreparedPRDiff) else None + if await self._prepare_chunked_prediction(model, prepared_diff): return if self.patches_diff: get_logger().debug("PR diff", diff=self.patches_diff) - self.prediction = await self._get_prediction(model) + prediction = await self._get_prediction(model) + self._load_valid_review_yaml(prediction) + self.prediction = prediction else: get_logger().warning(f"Empty diff for PR: {self.pr_url}") self.prediction = None - async def _prepare_chunked_prediction(self, model: str) -> bool: + async def _prepare_chunked_prediction(self, model: str, + prepared_diff: PreparedPRDiff | None = None) -> bool: """Review a too-large diff in chunks and merge the per-chunk verdicts. Returns False when chunking does not apply, leaving the single-call flow in place. """ - patches_diff_list, remaining_files_list = get_pr_multi_diffs( - self.git_provider, - self.token_handler, - model, - max_calls=get_settings().pr_reviewer.get("max_number_of_calls", 3), - add_line_numbers=True, - return_remaining_files=True) + patches_diff_list = getattr(self, "_chunked_patches_diff_list", None) + remaining_files_list = getattr(self, "_chunked_remaining_files_list", None) + if patches_diff_list is None: + multi_diff_kwargs = { + "max_calls": get_settings().pr_reviewer.get("max_number_of_calls", 3), + "add_line_numbers": True, + "return_remaining_files": True, + } + if prepared_diff is not None: + multi_diff_kwargs["prepared_diff"] = prepared_diff + patches_diff_list, remaining_files_list = get_pr_multi_diffs( + self.git_provider, + self.token_handler, + model, + **multi_diff_kwargs) + self._chunked_patches_diff_list = patches_diff_list + self._chunked_remaining_files_list = remaining_files_list if len(patches_diff_list) < 2: get_logger().info("Large-diff chunking produced a single chunk, reviewing the PR in one call") return False get_logger().info(f"Number of PR chunk calls: {len(patches_diff_list)}") get_logger().debug("PR diff chunks", artifact=patches_diff_list) + chunk_results = getattr(self, "_chunked_results", {}) + pending_indices = [index for index in range(len(patches_diff_list)) if index not in chunk_results] predictions = await asyncio.gather( - *[self._get_prediction(model, patches_diff) for patches_diff in patches_diff_list], + *[self._get_prediction(model, patches_diff_list[index]) for index in pending_indices], return_exceptions=True) - raw_predictions, chunk_outputs, chunk_errors = [], [], [] - for chunk_index, prediction in enumerate(predictions): + chunk_errors = [] + for chunk_index, prediction in zip(pending_indices, predictions, strict=True): if isinstance(prediction, Exception): chunk_errors.append(prediction) get_logger().warning(f"Failed to review chunk {chunk_index + 1}; retaining successful chunks", @@ -817,25 +845,28 @@ async def _prepare_chunked_prediction(self, model: str) -> bool: continue if isinstance(prediction, BaseException): raise prediction - data = self._load_review_yaml(prediction) - if not isinstance(data, dict) or not isinstance(data.get("review"), dict) or not data["review"]: - get_logger().warning(f"Failed to parse the review of chunk {chunk_index + 1}", - artifact={"data": data}) + try: + data = self._load_valid_review_yaml(prediction, source=f"review chunk {chunk_index + 1}") + except Exception as error: + chunk_errors.append(error) + get_logger().warning(f"Failed to parse review chunk {chunk_index + 1}; retrying it with fallback", + artifact={"error": error}) continue - raw_predictions.append(prediction) - chunk_outputs.append(data) + chunk_results[chunk_index] = (prediction, data) + self._chunked_results = chunk_results - if not chunk_outputs: + if len(chunk_results) < len(patches_diff_list): if chunk_errors: raise chunk_errors[0] - get_logger().warning("No chunk produced a parsable review, falling back to a single review call") - return False + raise ValueError("No valid review output was produced for one or more chunks") # the raw text is kept for logging only; the merged verdict is in self.prediction_data + raw_predictions = [chunk_results[index][0] for index in range(len(patches_diff_list))] + chunk_outputs = [chunk_results[index][1] for index in range(len(patches_diff_list))] self.prediction = "\n".join(raw_predictions) self.prediction_data = merge_review_chunks(chunk_outputs) self.review_chunk_count = len(patches_diff_list) - self.review_failed_chunk_count = len(patches_diff_list) - len(chunk_outputs) + self.review_failed_chunk_count = 0 self.remaining_files_list = remaining_files_list return True @@ -875,6 +906,14 @@ def _load_review_yaml(prediction: str) -> dict: "relevant_file:", "relevant_line:", "suggestion:"], first_key='review', last_key='security_concerns') + @classmethod + def _load_valid_review_yaml(cls, prediction: str, *, source: str = "model response") -> dict: + """Parse one prediction and require the minimum publishable review shape.""" + data = cls._load_review_yaml(prediction) + if not isinstance(data, dict) or not isinstance(data.get("review"), dict) or not data["review"]: + raise ValueError(f"{source} did not contain a non-empty review mapping") + return data + def _prepare_pr_review(self) -> str: """ Prepare the PR review by processing the AI prediction and generating a markdown-formatted text that summarizes @@ -1007,7 +1046,7 @@ def _prepare_pr_review(self) -> str: # Add custom labels from the review prediction (effort, security) self.set_review_labels(data) - if markdown_text == None or len(markdown_text) == 0: + if markdown_text is None or len(markdown_text) == 0: markdown_text = "" return markdown_text @@ -1166,7 +1205,9 @@ def _get_user_answers(self) -> Tuple[str, str]: question_str = "" answer_str = "" - if self.is_answer: + if self.is_answer and self.git_provider.is_supported("get_issue_comments"): + # __init__ already raises when is_answer is True and this is unsupported, so + # this repeats that check locally rather than relying on the caller for it. discussion_messages = self.git_provider.get_issue_comments() # providers return the comments oldest-first. PyGithub's PaginatedList reverses lazily, @@ -1186,29 +1227,6 @@ def _get_user_answers(self) -> Tuple[str, str]: return question_str, answer_str - def _get_previous_review_comment(self): - """ - Get the previous review comment if it exists. - """ - try: - if hasattr(self.git_provider, "get_previous_review"): - return self.git_provider.get_previous_review( - full=not self.incremental.is_incremental, - incremental=self.incremental.is_incremental, - ) - except Exception as e: - get_logger().exception(f"Failed to get previous review comment, error: {e}") - - def _remove_previous_review_comment(self, comment): - """ - Remove the previous review comment if it exists. - """ - try: - if comment: - self.git_provider.remove_comment(comment) - except Exception as e: - get_logger().exception(f"Failed to remove previous review comment, error: {e}") - def _can_run_incremental_review(self) -> bool: """ Checks if we can run incremental review according the various configurations and previous review. @@ -1233,12 +1251,17 @@ def _can_run_incremental_review(self) -> bool: num_commits_threshold = get_settings().pr_reviewer.minimal_commits_for_incremental_review not_enough_commits = num_new_commits < num_commits_threshold # checking if the commits are not too recent to start the review - recent_commits_threshold = datetime.datetime.now() - datetime.timedelta( + recent_commits_threshold = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) - datetime.timedelta( minutes=get_settings().pr_reviewer.minimal_minutes_for_incremental_review ) last_seen_commit_date = ( self.incremental.last_seen_commit.commit.author.date if self.incremental.last_seen_commit else None ) + # PyGithub returns timezone-aware UTC commit dates; the threshold below is a + # naive datetime. Normalize to naive UTC so the comparison cannot raise + # TypeError, matching how the GitLab and Azure providers emit commit dates. + if last_seen_commit_date is not None and last_seen_commit_date.tzinfo is not None: + last_seen_commit_date = last_seen_commit_date.astimezone(datetime.timezone.utc).replace(tzinfo=None) all_commits_too_recent = ( last_seen_commit_date > recent_commits_threshold if self.incremental.last_seen_commit else False ) diff --git a/pr_agent/tools/pr_similar_issue.py b/pr_agent/tools/pr_similar_issue.py index c952628384..1b8629c63d 100644 --- a/pr_agent/tools/pr_similar_issue.py +++ b/pr_agent/tools/pr_similar_issue.py @@ -67,10 +67,25 @@ def _qdrant_collection_name(base_name: str) -> str: return f"{base_name}-v2" +def _provider_supports_issue_indexing() -> bool: + """Whether the configured provider can back `/similar_issue`. + + The check is on the provider class rather than on the configured id, so a provider + registered through `register_git_provider()` is judged by the capability it declares. + An unresolvable configuration is reported as unsupported, which is what the tool's + `run()` already handles, rather than raised out of `__init__`. + """ + try: + provider_class = get_git_provider() + except ValueError: + return False + return provider_class.supports_issue_indexing() + + class PRSimilarIssue: def __init__(self, issue_url: str, ai_handler, args: list = None): self.issue_url = issue_url - self.supported = get_settings().config.git_provider == "github" + self.supported = _provider_supports_issue_indexing() if not self.supported: return @@ -88,39 +103,36 @@ def __init__(self, issue_url: str, ai_handler, args: list = None): if get_settings().pr_similar_issue.vectordb == "pinecone": try: import pinecone - except: - raise Exception("Please install 'pinecone' and 'pinecone_datasets' to use pinecone as vectordb") - # assuming pinecone api key and environment are set in secrets file + from pinecone import ServerlessSpec + except ImportError: + raise Exception("Please install the 'pinecone' package to use pinecone as vectordb") from None + # assuming pinecone api key, cloud and region are set in secrets file try: api_key = get_settings().pinecone.api_key - environment = get_settings().pinecone.environment + cloud = get_settings().pinecone.cloud + region = get_settings().pinecone.region except Exception: if not self.cli_mode: repo_name, original_issue_number = self.git_provider._parse_issue_url(self.issue_url.split('=')[-1]) issue_main = self.git_provider.repo_obj.get_issue(original_issue_number) - issue_main.create_comment("Please set pinecone api key and environment in secrets file") - raise Exception("Please set pinecone api key and environment in secrets file") + issue_main.create_comment("Please set pinecone api key, cloud and region in secrets file") + raise Exception("Please set pinecone api key, cloud and region in secrets file") + self.pc = pinecone.Pinecone(api_key=api_key) + self.pc_spec = ServerlessSpec(cloud=cloud, region=region) + self.pinecone_index = None # check if index exists, and if repo is already indexed run_from_scratch = False - if run_from_scratch: # for debugging - pinecone.init(api_key=api_key, environment=environment) - if index_name in pinecone.list_indexes(): - get_logger().info('Removing index...') - pinecone.delete_index(index_name) - get_logger().info('Done') - upsert = True - pinecone.init(api_key=api_key, environment=environment) - if index_name not in pinecone.list_indexes(): + if not self.pc.has_index(index_name): run_from_scratch = True upsert = False else: if get_settings().pr_similar_issue.force_update_dataset: upsert = True else: - pinecone_index = pinecone.Index(index_name=index_name) - res = pinecone_index.fetch([f"example_issue_{repo_name_for_index}"]).to_dict() + self.pinecone_index = self.pc.Index(name=index_name) + res = self.pinecone_index.fetch(ids=[f"example_issue_{repo_name_for_index}"]).to_dict() if res["vectors"]: upsert = False @@ -132,7 +144,7 @@ def __init__(self, issue_url: str, ai_handler, args: list = None): get_logger().info('Done') self._update_index_with_issues(issues, repo_name_for_index, upsert=upsert) else: # update index if needed - pinecone_index = pinecone.Index(index_name=index_name) + self.pinecone_index = self.pc.Index(name=index_name) issues_to_update = [] issues_paginated_list = repo_obj.get_issues(state='all') counter = 1 @@ -142,7 +154,7 @@ def __init__(self, issue_url: str, ai_handler, args: list = None): issue_str, comments, number = self._process_issue(issue) issue_key = f"issue_{number}" id = issue_key + "." + "issue" - res = pinecone_index.fetch([id]).to_dict() + res = self.pinecone_index.fetch(ids=[id]).to_dict() is_new_issue = True for vector in res["vectors"].values(): if vector['metadata']['repo'] == repo_name_for_index: @@ -310,7 +322,7 @@ def __init__(self, issue_url: str, ai_handler, args: list = None): async def run(self): if not self.supported: - message = "The /similar_issue tool is currently supported only for GitHub." + message = "The /similar_issue tool is not supported by the configured git provider." if get_settings().config.publish_output: try: from pr_agent.git_providers import get_git_provider_with_context @@ -337,10 +349,8 @@ async def run(self): score_list = [] if get_settings().pr_similar_issue.vectordb == "pinecone": - import pinecone - - pinecone_index = pinecone.Index(index_name=self.index_name) - res = pinecone_index.query(embeds[0], + pinecone_index = self.pc.Index(name=self.index_name) + res = pinecone_index.query(vector=embeds[0], top_k=5, filter={"repo": self.repo_name_for_index}, include_metadata=True).to_dict() @@ -451,8 +461,6 @@ def _process_issue(self, issue): def _update_index_with_issues(self, issues_list, repo_name_for_index, upsert=False): import pandas as pd - import pinecone - from pinecone_datasets import Dataset, DatasetMetadata get_logger().info('Processing issues...') corpus = Corpus() @@ -512,28 +520,26 @@ def _update_index_with_issues(self, issues_list, repo_name_for_index, upsert=Fal get_logger().info('Done') get_logger().info('Embedding...') - list_to_encode = list(df["text"].values) + list_to_encode = df["text"].to_list() embeds = _embed_with_fallback(list_to_encode) df["values"] = embeds - meta = DatasetMetadata.empty() - meta.dense_model.dimension = len(embeds[0]) - ds = Dataset.from_pandas(df, meta) get_logger().info('Done') - api_key = get_settings().pinecone.api_key - environment = get_settings().pinecone.environment + vectors = [ + (row["id"], row["values"], row["metadata"]) + for row in df.to_dict(orient="records") + ] if not upsert: get_logger().info('Creating index from scratch...') - ds.to_pinecone_index(self.index_name, api_key=api_key, environment=environment) - time.sleep(15) # wait for pinecone to finalize indexing before querying - else: - get_logger().info('Upserting index...') - namespace = "" - batch_size: int = 100 - concurrency: int = 10 - pinecone.init(api_key=api_key, environment=environment) - ds._upsert_to_index(self.index_name, namespace, batch_size, concurrency) - time.sleep(5) # wait for pinecone to finalize upserting before querying + self.pc.create_index(name=self.index_name, dimension=len(embeds[0]), metric="cosine", spec=self.pc_spec, + timeout=120) + get_logger().info('Upserting index...') + self.pinecone_index = self.pc.Index(name=self.index_name) + self.pinecone_index.upsert(vectors=vectors, + namespace="", + batch_size=100, + max_concurrency=10) + time.sleep(5) # wait for pinecone to finalize upserting before querying get_logger().info('Done') def _update_table_with_issues(self, issues_list, repo_name_for_index, ingest=False): @@ -598,7 +604,7 @@ def _update_table_with_issues(self, issues_list, repo_name_for_index, ingest=Fal get_logger().info('Done') get_logger().info('Embedding...') - list_to_encode = list(df["text"].values) + list_to_encode = df["text"].to_list() embeds = _embed_with_fallback(list_to_encode) df["vector"] = embeds get_logger().info('Done') @@ -685,7 +691,7 @@ def _update_qdrant_with_issues(self, issues_list, repo_name_for_index, ingest=Fa get_logger().info('Done') get_logger().info('Embedding...') - list_to_encode = list(df["text"].values) + list_to_encode = df["text"].to_list() embeds = _embed_with_fallback(list_to_encode) df["vector"] = embeds get_logger().info('Done') diff --git a/pr_agent/tools/pr_update_changelog.py b/pr_agent/tools/pr_update_changelog.py index 016b659856..d9126e3d1a 100644 --- a/pr_agent/tools/pr_update_changelog.py +++ b/pr_agent/tools/pr_update_changelog.py @@ -167,6 +167,11 @@ def _prepare_changelog_update(self) -> Tuple[str, str]: return new_file_content, answer def _push_changelog_update(self, new_file_content, answer): + if not self.git_provider.is_supported("push_code"): + # Its only caller already gates on self.commit_changelog, which is False + # whenever this capability is missing; kept local so the guard holds even + # if a future caller reaches this method some other way. + return if get_settings().pr_update_changelog.get("skip_ci_on_push", True): commit_message = "[skip ci] Update CHANGELOG.md" else: @@ -180,7 +185,7 @@ def _push_changelog_update(self, new_file_content, answer): sleep(5) # wait for the file to be updated try: - if get_settings().config.git_provider == "github": + if self.git_provider.supports_changelog_update_review(): last_commit_id = list(self.git_provider.pr.get_commits())[-1] d = dict( body="CHANGELOG.md update", diff --git a/pr_agent/tools/ticket_pr_compliance_check.py b/pr_agent/tools/ticket_pr_compliance_check.py index 797439aeec..eabc92aad3 100644 --- a/pr_agent/tools/ticket_pr_compliance_check.py +++ b/pr_agent/tools/ticket_pr_compliance_check.py @@ -1,3 +1,4 @@ +import asyncio import copy import math import re @@ -5,17 +6,20 @@ from urllib.parse import urlparse import aiohttp +from atlassian import Jira from pr_agent.algo.pr_processing import OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD from pr_agent.algo.token_handler import TokenHandler from pr_agent.algo.utils import get_max_tokens from pr_agent.config_loader import get_settings -from pr_agent.git_providers import AzureDevopsProvider, GithubProvider, GitLabProvider +from pr_agent.git_providers.git_provider import GitProvider from pr_agent.log import get_logger # Compile the regex pattern once, outside the function GITHUB_TICKET_PATTERN = re.compile( - r'(https://github[^/]+/[^/]+/[^/]+/issues/\d+)|(\b(\w+)/(\w+)#(\d+)\b)|(#\d+)' + r'(https://github[^/]+/[^/]+/[^/]+/issues/\d+)' + r'|((?" in https://.atlassian.net). Only Jira Cloud is +# supported: the base URL is built from this name rather than taken as a free-form URL, so +# repo-controlled config cannot redirect the authenticated request (and the token) to an +# arbitrary host. The name must be a single DNS label (letters, digits, hyphens) so it +# cannot contain '.', '/', ':', '@' etc. that would let it escape *.atlassian.net. +JIRA_SITE_PATTERN = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$", re.IGNORECASE) + +# A Jira project key as configured in jira.project_keys: the "" in PROJECT-123. +# Bounded like the prefix in find_jira_tickets (2-10 letters) so an entry can only ever be +# a plain label, never a URL or a full ticket key. Case-sensitive: Jira project keys are +# upper case, so a configured "proj" is a typo rather than a spelling variant, and +# accepting it silently would widen the lookup to a project the repo never named. Keys +# found in the PR text are still matched case-insensitively and upper-cased +# (find_jira_tickets) before they are compared. +JIRA_PROJECT_KEY_PATTERN = re.compile(r"^[A-Z]{2,10}$") + + +def _jira_cloud_base_url(): + """ + Return the Jira Cloud base URL built from the configured site name, or None if the + site is missing or invalid. Building the URL from a validated site name (rather than + accepting a free-form base URL from settings) ensures the authenticated request can + only ever go to https://.atlassian.net, even if settings were overridden by + untrusted repo configuration. + """ + site = get_settings().get("JIRA.JIRA_SITE", None) + if not site: + return None + site = str(site).strip() + if not JIRA_SITE_PATTERN.match(site): + get_logger().warning( + f"Invalid jira_site '{site}'; expected a Jira Cloud site name like 'mycompany' " + f"(the '' in https://.atlassian.net). Skipping Jira ticket lookup.") + return None + return f"https://{site}.atlassian.net" + + +def _jira_project_keys(): + """ + Return the configured jira.project_keys allowlist as a set of project keys, or None + when nothing was supplied (look up every key found). + + Only three values count as "nothing was supplied": the option missing entirely, the + shipped empty list, and a top-level string that is blank, the last because + `jira__project_keys=""` is how a shell spells an unset environment variable. + + Everything else is a supplied value and is judged on its entries. Entries that are not + plain upper-case project keys, including a blank one (so `[""]` is a malformed + override rather than a second spelling of the default), are ignored with a warning, + and a supplied value left with no valid entry yields an empty set, which drops every + key. So does a value that is not a list at all (a boolean or number from a YAML or CLI + override). A typo fails closed instead of silently widening the lookup to every + key-shaped string. Accepts a list or a comma-separated string, the latter for + environment-variable overrides (jira__project_keys="PROJ,OPS"). + """ + configured = get_settings().get("JIRA.PROJECT_KEYS", None) + if configured is None: + return None + if isinstance(configured, str): + if not configured.strip(): + return None + configured = configured.split(",") + elif not isinstance(configured, (list, tuple)): + # A YAML/CLI override such as `project_keys=true` or `=false` is neither + # "unset" nor a list; fail closed rather than guess what it meant. + get_logger().warning( + f"jira.project_keys must be a list of project keys, got {type(configured).__name__}; " + "skipping Jira ticket lookup until it is fixed") + return set() + if not configured: + # The shipped default. An empty container carries no entry to be wrong about. + return None + allowed = set() + for index, item in enumerate(configured): + # Only strings are candidates: a TOML/YAML boolean or null in the list must not + # be stringified into a key-shaped label ("TRUE", "NONE") that then filters. + key = item.strip() if isinstance(item, str) else item + if not isinstance(key, str) or not JIRA_PROJECT_KEY_PATTERN.match(key): + # The entry is named by position and type, never quoted back. Its content is + # configuration this tool does not control: it can be long, hold a secret + # pasted into the wrong setting, or carry newlines that would forge log + # lines. The position is what an operator needs to find and fix it. + get_logger().warning( + f"Ignoring invalid jira.project_keys entry at index {index} " + f"(type {type(item).__name__}); expected a plain upper-case project key " + "like 'PROJ'") + continue + allowed.add(key) + if not allowed: + get_logger().warning( + "jira.project_keys has no valid entry; skipping Jira ticket lookup until it is fixed") + return allowed + def find_jira_tickets(text): - # Regular expression patterns for JIRA tickets + # Regular expression patterns for JIRA tickets. Matching is case-insensitive so + # lowercased branch names (e.g. bugfix/abc-123-description) are detected; keys are + # normalized to upper case to match Jira's canonical form. patterns = [ r'\b[A-Z]{2,10}-\d{1,7}\b', # Standard JIRA ticket format (e.g., PROJ-123) r'(?:https?://[^\s/]+/browse/)?([A-Z]{2,10}-\d{1,7})\b' # JIRA URL or just the ticket ] - tickets = set() + # Preserve first-seen order while de-duplicating, so the MAX_TICKETS cap applied + # later is deterministic across runs (a plain set would fetch arbitrary tickets). + seen = set() + tickets = [] for pattern in patterns: - matches = re.findall(pattern, text) + matches = re.findall(pattern, text, flags=re.IGNORECASE) for match in matches: if isinstance(match, tuple): # If it's a tuple (from the URL pattern), take the last non-empty group @@ -42,9 +161,187 @@ def find_jira_tickets(text): else: ticket = match if ticket: - tickets.add(ticket) + ticket = ticket.upper() + if ticket not in seen: + seen.add(ticket) + tickets.append(ticket) + + return tickets + + +def _get_jira_client(): + """ + Build a Jira Cloud client from the [jira] settings. Returns None if Jira is not + configured. Only Jira Cloud is supported: the base URL is derived from a validated + site name (jira_site) so untrusted repo configuration cannot redirect the + authenticated request to an arbitrary host. Cloud authenticates with the account + email + API token. The REST API version is pinned via JIRA_API_VERSION (see its + definition for why). + """ + site = get_settings().get("JIRA.JIRA_SITE", None) + api_email = get_settings().get("JIRA.JIRA_API_EMAIL", None) + api_token = get_settings().get("JIRA.JIRA_API_TOKEN", None) + base_url = _jira_cloud_base_url() # None if site is missing or invalid (already warned if invalid) + if not (base_url and api_email and api_token): + # Warn when Jira is partially configured: some [jira] value is set but the required + # site + email + token are incomplete, which is likely a mistake. Stay silent when + # nothing is set (Jira simply not in use), and don't double-warn when the site was + # set but invalid (_jira_cloud_base_url already warned about that). + site_set_but_invalid = site and not base_url + if any([site, api_email, api_token]) and not site_set_but_invalid: + missing = [ + name for name, value in ( + ("jira_site", site), + ("jira_api_email", api_email), + ("jira_api_token", api_token), + ) if not value + ] + get_logger().warning( + f"Jira is partially configured; skipping Jira ticket lookup. Missing: {', '.join(missing)}") + return None + try: + return Jira(url=base_url, username=api_email, password=api_token, api_version=JIRA_API_VERSION) + except Exception as e: + get_logger().error(f"Failed to initialize Jira client: {e}", + artifact={"traceback": traceback.format_exc()}) + return None + + +def extract_jira_tickets(text, max_characters=MAX_TICKET_CHARACTERS, max_tickets=MAX_TICKETS): + """ + Find Jira ticket keys in the given text and fetch their content. Returns a list of + ticket dicts in the same shape used by the rest of the ticket-analysis flow. Returns + an empty list when no keys are found or when Jira is not configured. + + Candidates are tried in first-seen order until max_tickets resolve or + MAX_JIRA_FETCH_ATTEMPTS lookups have been made. Counting resolved tickets rather than + truncating the candidate list keeps key-shaped noise ("SHA-256", "UTF-8") from + displacing a real ticket that appears later in the text. + + When jira.project_keys is configured, candidates whose prefix isn't in that allowlist + are dropped before any lookup, so they don't consume a fetch attempt. + """ + # Look for keys before building a client: most PRs have none, and building the + # client first would do needless work (and log a noisy init failure if Jira is + # misconfigured) even when there is nothing to fetch. + keys = find_jira_tickets(text or "") + if not keys: + return [] + + # Optional allowlist: key-shaped false positives like "SHA-256" or "UTF-8" match the + # regex but aren't real Jira keys. Filtering by known project prefixes here means they + # never consume one of the MAX_JIRA_FETCH_ATTEMPTS lookups below. Empty keeps today's + # behaviour; a supplied but invalid allowlist drops every key (see _jira_project_keys). + allowed_projects = _jira_project_keys() + if allowed_projects is not None: + # Partitioned in one pass: the candidate list is not capped until lookups begin, + # so a PR carrying many distinct key-shaped tokens would otherwise pay a + # membership scan of the skipped list for every one of them. + kept, skipped = [], [] + for key in keys: + target = kept if key.split("-", 1)[0] in allowed_projects else skipped + target.append(key) + if skipped: + get_logger().debug( + f"Skipping Jira lookup for keys outside jira.project_keys: {', '.join(skipped)}") + keys = kept + if not keys: + return [] + + jira_client = _get_jira_client() + if jira_client is None: + return [] + + base_url = _jira_cloud_base_url() or "" + # Custom field that holds acceptance criteria / requirements. The field id is + # instance-specific (e.g. "customfield_10127"), so it must be configured; empty + # means no requirements are extracted. + requirements_field = get_settings().get("JIRA.JIRA_REQUIREMENTS_FIELD", "") or "" + tickets_content = [] + for attempt, key in enumerate(keys, start=1): + if len(tickets_content) >= max_tickets: + get_logger().info(f"Reached the cap of {max_tickets} Jira tickets; skipping remaining keys") + break + if attempt > MAX_JIRA_FETCH_ATTEMPTS: + get_logger().info( + f"Stopped after {MAX_JIRA_FETCH_ATTEMPTS} Jira lookups; skipping remaining keys") + break + try: + issue = jira_client.issue(key) + except Exception as e: + get_logger().warning(f"Failed to fetch Jira ticket {key}: {e}") + continue + if not issue: + continue + + fields = issue.get("fields", {}) or {} + # The client is pinned to REST v2 (see _get_jira_client), which returns + # description and rich-text custom fields as plain wiki-markup strings. The + # isinstance guards below defend against anything non-string (e.g. v3 ADF dicts). + body = fields.get("description") or "" + if not isinstance(body, str): + body = "" + if len(body) > max_characters: + body = body[:max_characters] + "..." + + requirements = "" + if requirements_field: + requirements = fields.get(requirements_field) or "" + if not isinstance(requirements, str): + requirements = "" + if len(requirements) > max_characters: + requirements = requirements[:max_characters] + "..." + + labels = fields.get("labels", []) or [] + tickets_content.append({ + "ticket_id": key, + "ticket_url": f"{base_url}/browse/{key}" if base_url else "", + "title": fields.get("summary", ""), + "body": body, + "requirements": requirements, + "labels": ", ".join(labels), + }) + return tickets_content - return list(tickets) + +def _get_pr_title(git_provider) -> str: + """Return the PR/MR title across providers (GitHub/Bitbucket use .pr, GitLab .mr).""" + for attr in ("pr", "mr"): + obj = getattr(git_provider, attr, None) + title = getattr(obj, "title", None) + if title: + return title + return "" + + +def add_jira_tickets(git_provider, tickets_content): + """ + Provider-agnostic Jira lookup. Scans the PR title, description and branch name for + Jira ticket keys and appends any found tickets to tickets_content (de-duplicated by + ticket_url). No-op when Jira is not configured. Works for any git provider, since it + only relies on get_user_description() and get_pr_branch(). + + MAX_TICKETS is the overall per-PR cap, so any provider-native tickets already in + tickets_content count against it: only the remaining budget is offered to the Jira + lookup, keeping the existing tickets first. + """ + try: + if len(tickets_content) >= MAX_TICKETS: + return tickets_content + jira_context = "\n".join(filter(None, [ + _get_pr_title(git_provider), + git_provider.get_user_description() or "", + git_provider.get_pr_branch() or "", + ])) + existing_urls = {t.get("ticket_url") for t in tickets_content} + remaining = MAX_TICKETS - len(tickets_content) + for jira_ticket in extract_jira_tickets(jira_context, MAX_TICKET_CHARACTERS, remaining): + if jira_ticket.get("ticket_url") not in existing_urls: + tickets_content.append(jira_ticket) + except Exception as e: + get_logger().error(f"Error extracting Jira tickets: {e}", + artifact={"traceback": traceback.format_exc()}) + return tickets_content _ASANA_TASK_URL_PATTERN = re.compile( @@ -224,26 +521,39 @@ async def _fetch_asana_ticket_contents( timeout = aiohttp.ClientTimeout(total=_get_asana_request_timeout()) headers = {"Authorization": f"Bearer {api_token.strip()}"} - tickets_content = [] async with aiohttp.ClientSession(headers=headers, timeout=timeout) as session: # Bound attempts as well as successful results so invalid references cannot # multiply request latency, rate-limit usage, or warning logs. + selected_tickets = [] for ticket_url in ticket_urls[:max_tickets]: - if len(tickets_content) >= max_tickets: - break - task_gid = None try: task_gid = _get_asana_task_gid(ticket_url) - ticket_content = await _fetch_asana_ticket_content( - session, - ticket_url, - max_body_characters, - ) except Exception as e: - task_label = task_gid or "invalid reference" - get_logger().warning(f"Failed to fetch Asana task {task_label}: {e}") + get_logger().warning(f"Failed to fetch Asana task invalid reference: {e}") continue - tickets_content.append(ticket_content) + selected_tickets.append((ticket_url, task_gid)) + + async def fetch_ticket(ticket_url): + try: + return await _fetch_asana_ticket_content(session, ticket_url, max_body_characters), None + except Exception as e: + return None, e + + tasks = [asyncio.create_task(fetch_ticket(ticket_url)) for ticket_url, _task_gid in selected_tickets] + try: + results = await asyncio.gather(*tasks) + except BaseException: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + tickets_content = [] + for (_ticket_url, task_gid), (ticket_content, error) in zip(selected_tickets, results, strict=True): + if error is not None: + get_logger().warning(f"Failed to fetch Asana task {task_gid}: {error}") + continue + tickets_content.append(ticket_content) return tickets_content @@ -460,8 +770,23 @@ def _get_repo_obj_for_ticket(git_provider, ticket_url, repo_name, repo_obj_cache return repo_obj +def _provider_supports(git_provider, capability: str) -> bool: + """Read an opt-in ticket capability from a provider. + + Objects outside the GitProvider hierarchy (test doubles, minimal adapters) may not + define the method at all; absence means the capability is not supported, which keeps + the previous behaviour for providers that matched none of the concrete classes. + + A permissive double such as a bare MagicMock answers every valid capability truthily and + so takes the first branch. Unknown capabilities raise to expose misspelled names. + """ + if not hasattr(GitProvider, capability): + raise AttributeError(f"unknown provider capability: {capability!r}") + check = getattr(git_provider, capability, None) + return bool(check()) if callable(check) else False + + async def extract_tickets(git_provider): - MAX_TICKET_CHARACTERS = 10000 try: user_description = _get_user_description_for_asana(git_provider) asana_ticket_urls = find_asana_tickets(user_description) @@ -475,7 +800,7 @@ async def extract_tickets(git_provider): get_logger().warning(f"Failed to initialize Asana task fetching: {e}") asana_tickets_content = [] - if isinstance(git_provider, GithubProvider): + if _provider_supports(git_provider, "supports_issue_url_tickets"): description_tickets = extract_ticket_links_from_pr_description( user_description, git_provider.repo, git_provider.base_url_html ) @@ -570,9 +895,11 @@ async def extract_tickets(git_provider): }) tickets_content.extend(asana_tickets_content) + # Provider-agnostic Jira lookup (see add_jira_tickets); no-op when Jira is unconfigured. + add_jira_tickets(git_provider, tickets_content) return tickets_content - elif isinstance(git_provider, GitLabProvider): + elif _provider_supports(git_provider, "supports_issue_reference_tickets"): references = extract_gitlab_ticket_references( user_description, git_provider.id_project, @@ -581,7 +908,8 @@ async def extract_tickets(git_provider): tickets_content = [] for project_path, issue_iid in references: try: - project = git_provider.gl.projects.get(project_path) + # Only the issue manager is needed; avoid fetching unused project metadata. + project = git_provider.gl.projects.get(project_path, lazy=True) issue = project.issues.get(issue_iid) except Exception as e: get_logger().error( @@ -605,9 +933,11 @@ async def extract_tickets(git_provider): ) tickets_content.extend(asana_tickets_content) + # Provider-agnostic Jira lookup (see add_jira_tickets); no-op when Jira is unconfigured. + add_jira_tickets(git_provider, tickets_content) return tickets_content - elif isinstance(git_provider, AzureDevopsProvider): + elif _provider_supports(git_provider, "supports_linked_work_item_tickets"): tickets_info = git_provider.get_linked_work_items() tickets_content = [] for ticket in tickets_info: @@ -616,13 +946,21 @@ async def extract_tickets(git_provider): if len(ticket_body_str) > MAX_TICKET_CHARACTERS: ticket_body_str = ticket_body_str[:MAX_TICKET_CHARACTERS] + "..." + # Cap acceptance criteria like the body, so a large work-item field + # can't push an unbounded blob into the review prompt. + requirements_str = ticket.get("acceptance_criteria", "") or "" + if not isinstance(requirements_str, str): + requirements_str = "" + if len(requirements_str) > MAX_TICKET_CHARACTERS: + requirements_str = requirements_str[:MAX_TICKET_CHARACTERS] + "..." + tickets_content.append( { "ticket_id": ticket.get("id"), "ticket_url": ticket.get("url"), "title": ticket.get("title"), "body": ticket_body_str, - "requirements": ticket.get("acceptance_criteria", ""), + "requirements": requirements_str, "labels": ", ".join(ticket.get("labels", [])), } ) @@ -634,10 +972,18 @@ async def extract_tickets(git_provider): # Preserve the existing Azure work-item result set. The independently bounded # Asana results add context without imposing a new cap on Azure's established behaviour. tickets_content.extend(asana_tickets_content) + # Provider-agnostic Jira lookup (see add_jira_tickets); no-op when Jira is unconfigured. + add_jira_tickets(git_provider, tickets_content) return tickets_content - if asana_ticket_urls: - return asana_tickets_content + # Providers with no ticket integration of their own still reach Jira: keys are + # usually referenced in the PR title, description or branch name rather than by a + # provider-native link. The Asana results seed the list so they count against + # MAX_TICKETS, as on the provider-specific paths. Returning None when nothing at + # all was found keeps the "provider unsupported" contract intact. + tickets_content = add_jira_tickets(git_provider, list(asana_tickets_content)) + if asana_ticket_urls or tickets_content: + return tickets_content except Exception as e: get_logger().error(f"Error extracting tickets error= {e}", @@ -673,7 +1019,3 @@ async def extract_and_cache_pr_tickets(git_provider, vars): else: get_logger().info("Using cached tickets", artifact={"tickets": related_tickets}) vars['related_tickets'] = related_tickets - - -def check_tickets_relevancy(): - return True diff --git a/pyproject.toml b/pyproject.toml index 6fba1427be..4b959637a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,23 +36,25 @@ dependencies = [ "azure-devops==7.1.0b4", "azure-identity==1.25.0", "boto3==1.43.78", - "certifi==2024.8.30", + "certifi>=2026.7.22", "dynaconf>=3.2.13,<4", "fastapi>=0.141.1,<1", "GitPython>=3.1.59,<4", "starlette>=1.6.0,<2", "google-cloud-aiplatform==1.154.0", # google-cloud-aiplatform==1.154.0 requires storage <4,>=1.32.0 on py<3.13 and <4,>=3.10.0 on py>=3.13. - # Pin per-Python (exact, like the rest of this list) so existing installs are unchanged and py3.13 resolves. See #2480. + # Pin per-Python so existing installs are unchanged and py3.13 resolves. See #2480. "google-cloud-storage==2.10.0; python_version < '3.13'", "google-cloud-storage==3.12.0; python_version >= '3.13'", - "protobuf==6.33.6", # grpcio-status 1.80.0 forces >=6.31.1,<7; pinned to lock the resolved transitive version + # Keep protobuf explicit to lock the transitive version resolved with the gRPC dependency stack. + # Revalidate compatibility when updating grpcio-status or related Google/gRPC dependencies. + "protobuf==6.33.6", "Jinja2==3.1.6", "litellm==1.100.0", "loguru==0.7.2", "msrest==0.7.1", "openai>=1.55.3", - "PyGithub==1.59.*", + "PyGithub>=2.10,<3", "PyJWT>=2.13.0,<3", "PyYAML==6.0.1", "python-gitlab==8.3.0", @@ -78,6 +80,9 @@ dependencies = [ "opentelemetry-api==1.43.0", "opentelemetry-sdk==1.43.0", "opentelemetry-exporter-otlp-proto-http==1.43.0", + # Prometheus scrape endpoint for [otel] command telemetry (native /metrics on + # the webhook apps). Multiprocess-capable; see pr_agent/telemetry/prometheus.py. + "prometheus-client==0.21.1", ] [project.urls] @@ -89,6 +94,9 @@ dependencies = [ [tool.setuptools] include-package-data = true +[tool.setuptools.package-data] +pr_agent = ["py.typed"] + [tool.setuptools.packages.find] where = ["."] include = [ @@ -113,13 +121,11 @@ pr-agent = "pr_agent.cli:run" [dependency-groups] # 'similar issue' tool. Install with: uv sync --group similar-issue. -# Keep this a dependency group rather than a published extra: pinecone-datasets -# resolves from a git fork, which wheel metadata cannot reference (PyPI rejects -# direct URL deps), so a pip-installed extra would silently pull the wrong package. +# Keep this a dependency group rather than a published extra so the vector +# database drivers stay out of the base install. similar-issue = [ "pandas", # imported directly by the qdrant indexing path - "pinecone-client", - "pinecone-datasets @ git+https://github.com/mrT23/pinecone-datasets.git@main", + "pinecone==10.0.0", "lancedb==0.38.0", "qdrant-client==1.15.1", ] @@ -132,11 +138,16 @@ dev = [ "ruff==0.16.5", ] +[tool.uv] +# Keep this exact pin in sync with docker/Dockerfile and docker/Dockerfile.lambda +# so local, CI, release, and container builds use the same uv version. +required-version = "==0.12.10" + [tool.ruff] line-length = 120 lint.select = [ - "E", # Pyflakes + "E", # pycodestyle "F", # Pyflakes "B", # flake8-bugbear "I", # isort @@ -154,19 +165,14 @@ lint.unfixable = [ lint.exclude = ["api/code_completions"] -# B008 is a long-standing ignore. Everything below it is a pre-existing +# Everything below is a pre-existing # violation deferred so ruff runs green on adoption — fix the code and drop # the entry. TODO: burn this list down. lint.ignore = [ - "B008", # function call in argument defaults "E402", # import not at top of file "E501", # line too long - "E711", # comparison to None "E722", # bare `except` - "F811", # redefinition of unused name "F841", # unused local variable - "B006", # mutable argument default - "B007", # unused loop control variable "B011", # `assert False` "B017", # assertRaises(Exception) "B023", # function definition uses loop variable @@ -177,10 +183,6 @@ lint.ignore = [ [tool.ruff.lint.per-file-ignores] # Ignore import-position and unused-import in __init__.py (re-export modules). "__init__.py" = ["E402", "F401"] -[tool.bandit] -exclude_dirs = ["tests"] -skips = ["B101"] -tests = [] [tool.pytest.ini_options] asyncio_mode = "auto" diff --git a/scripts/generate_config_reference.py b/scripts/generate_config_reference.py new file mode 100644 index 0000000000..f66229b52e --- /dev/null +++ b/scripts/generate_config_reference.py @@ -0,0 +1,232 @@ +"""Generate docs/docs/usage-guide/configuration_reference.md from configuration.toml. + +The generated page lists every active option in pr_agent/settings/configuration.toml, +grouped by section, rendering each default value and its inline comment so the docs +stay in sync with the authoritative source. Run from the repo root: + + python scripts/generate_config_reference.py [output_path] + +The output path defaults to docs/docs/usage-guide/configuration_reference.md. +""" + +from __future__ import annotations + +import re +import sys +import tomllib +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +CONFIG_TOML = ROOT / "pr_agent/settings/configuration.toml" +DEFAULT_OUTPUT = ROOT / "docs/docs/usage-guide/configuration_reference.md" + +SOURCE_URL = "https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml" + +PAGE_HEADER = f"""# Configuration Reference + +> This page is **auto-generated** and should not be edited by hand. +> Regenerate it from the [TOML source]({SOURCE_URL}) with: +> +> ```bash +> python scripts/generate_config_reference.py +> ``` + +Every configuration option PR-Agent supports, grouped by section. The [configuration.toml]({SOURCE_URL}) +file is the single source of truth for defaults and inline comments; this page renders the same +list for easy searching and linking. + +Rows with an empty **Description** are keys whose TOML entry carries no explanatory comment yet. +They are listed deliberately rather than hidden, so the gaps double as the documentation +to-do list. +""" + +TABLE_HEADER = "| Key | Default | Description |\n| --- | --- | --- |" + +# Comment lines shorter than this without trailing sentence punctuation are treated as +# group labels (e.g. "# models") rather than descriptions of the following key. +LABEL_MAX_LEN = 60 + +KEY_RE = re.compile(r"^(?P[a-zA-Z0-9_]+)\s*=\s*(?P.*)$") +SECTION_RE = re.compile(r"^\[(?P[^\]]+)\](?:\s*#+\s*(?P.+?)\s*#*)?$") + + +def _is_assignment(text: str) -> bool: + return bool(re.match(r"^[a-zA-Z0-9_\"']+\s*=", text)) + + +def _classify_comment(text: str) -> str: + stripped = text.strip().lstrip("#").strip() + if not stripped: + return "empty" + if _is_assignment(stripped): + return "example" + if len(stripped) <= LABEL_MAX_LEN and not re.search(r"[.:!?]", stripped): + return "label" + return "desc" + + +def _probe_value(buf: str) -> tuple: + """Return (parsed_value, value_text, inline_comment) or (None, None, None).""" + idx = buf.find("#") + if idx != -1 and buf[:idx].rstrip(): + candidate = buf[:idx].rstrip() + try: + return tomllib.loads(f"k = {candidate}")["k"], candidate, buf[idx + 1 :].strip() + except tomllib.TOMLDecodeError: + pass + try: + return tomllib.loads(f"k = {buf}")["k"], buf, None + except tomllib.TOMLDecodeError: + return None, None, None + + +def _render_value(value: object) -> str: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)): + return str(value) + if isinstance(value, str): + if "\n" in value: + lines = [line.strip() for line in value.strip().splitlines() if line.strip()] + first = lines[0] if lines else "" + text = first + (" ..." if len(lines) > 1 else "") + return '"""' + text + '"""' + return '"' + value.replace("\\", "\\\\").replace('"', '\\"') + '"' + if isinstance(value, list): + return "[" + ", ".join(_render_value(item) for item in value) + "]" + if isinstance(value, dict): + items = ", ".join(f"{_render_value(key)} = {_render_value(item)}" for key, item in value.items()) + return "{" + items + "}" + return str(value) + + +def _escape_cell(text: str) -> str: + return text.replace("|", "\\|").replace("\n", " ") + + +def load_sections(text: str) -> list: + sections = [] + current = None + pending_comments = [] + + lines = text.splitlines() + i = 0 + while i < len(lines): + line = lines[i].strip() + + section_match = SECTION_RE.match(line) + if section_match: + current = {"name": section_match.group("name"), "tag": (section_match.group("tag") or "").strip()} + current["keys"] = [] + sections.append(current) + pending_comments = [] + i += 1 + continue + + if not line: + i += 1 + continue + + if line.startswith("#"): + if current is not None: + kind = _classify_comment(line) + if kind in ("example", "empty"): + pending_comments = [] + else: + pending_comments.append(line.lstrip("#").strip()) + i += 1 + continue + + key_match = KEY_RE.match(line) + if key_match is None or current is None: + i += 1 + continue + + buf = key_match.group("rest") + value, value_text, inline_comment = _probe_value(buf) + i += 1 + while value is None and i < len(lines): + buf += "\n" + lines[i] + value, value_text, inline_comment = _probe_value(buf) + i += 1 + if value is None: + value_text, inline_comment = buf, None + + desc = inline_comment + if desc is None: + desc_parts = [c for c in pending_comments if _classify_comment(c) == "desc"] + desc = " ".join(desc_parts) or None + label_parts = [c for c in pending_comments if _classify_comment(c) == "label"] + label = label_parts[-1] if label_parts else None + pending_comments = [] + + current["keys"].append( + { + "key": key_match.group("key"), + "value": _render_value(value) if value is not None else value_text, + "desc": desc, + "label": label, + } + ) + + return sections + + +def render_page(sections: list) -> str: + out = [PAGE_HEADER.rstrip(), ""] + + for section in sections: + keys = section["keys"] + heading = f"## `[{section['name']}]`" + if section["tag"]: + heading += f" — {section['tag']}" + out.append(heading) + out.append("") + if not keys: + out.append("_This section only documents commented-out examples; see the [TOML source]" + f"({SOURCE_URL}) for details._") + out.append("") + continue + + table_open = False + pending_label = None + for key in keys: + if key["label"]: + if table_open: + table_open = False + pending_label = "**" + key["label"] + "**" + if not table_open: + if pending_label: + out.append(pending_label) + out.append("") + pending_label = None + out.append(TABLE_HEADER) + table_open = True + desc = _escape_cell(key["desc"]) if key["desc"] else "" + value = _escape_cell(key["value"]) + out.append(f"| `{key['key']}` | {value if value.strip() else '`""`'} | {desc} |") + if table_open: + out.append("") + out.append("") + + return "\n".join(out).rstrip("\n") + "\n" + + +def main() -> int: + output = Path(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_OUTPUT + text = CONFIG_TOML.read_text(encoding="utf-8") + sections = load_sections(text) + + rendered_keys = sum(len(section["keys"]) for section in sections) + expected_keys = len(re.findall(r"^[a-zA-Z0-9_]+\s*=", text, re.M)) + if rendered_keys != expected_keys: + print(f"ERROR: rendered {rendered_keys} keys, expected {expected_keys}", file=sys.stderr) + return 1 + + output.write_text(render_page(sections), encoding="utf-8") + print(f"Wrote {rendered_keys} keys across {len(sections)} sections to {output}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/health_test/test_github_smoke.py b/tests/health_test/test_github_smoke.py new file mode 100644 index 0000000000..b0d3963320 --- /dev/null +++ b/tests/health_test/test_github_smoke.py @@ -0,0 +1,116 @@ +import asyncio +import copy +import os + +from starlette_context import context, request_cycle_context + +from pr_agent.agent.pr_agent import PRAgent +from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler +from pr_agent.config_loader import get_settings, global_settings + +# The default target is linked to #2934, so the smoke run +# exercises the ticket extraction path against live data. Override with +# TEST_PR_URL; assertions specific to the default target are then skipped. +DEFAULT_PR_URL = 'https://github.com/The-PR-Agent/pr-agent/pull/2940' + +DESCRIBE_STUB_YAML = """\ +title: | + fix: guard against empty commit list in Azure DevOps get_latest_commit_url +type: Enhancement +description: | + Synthetic PR description produced by the CI smoke stub so the run needs no model key. +pr_files: + - filename: | + pr_agent/git_providers/azuredevops_provider.py + changes_title: | + Guard the empty commit list + changes_summary: | + Return an empty string instead of raising IndexError when the commit list is empty + label: | + Enhancement + - filename: | + tests/unittest/test_azure_devops_comment.py + changes_title: | + Add empty-commit regression coverage + changes_summary: | + Cover get_latest_commit_url with an empty commit list + label: | + Tests +""" + + +class StubDescribeHandler(BaseAiHandler): + """Deterministic AI handler for the GitHub smoke test. + + Returns fixed YAML valid for the /describe pipeline, so the run exercises + provider auth, diff building, language sorting, ticket extraction and prompt + rendering end to end without any model key. + """ + + def __init__(self): + self.main_pr_language = "Python" + self.calls = [] + + @property + def deployment_id(self): + return "smoke-stub" + + async def chat_completion(self, model: str, system: str, user: str, + temperature: float = 0.2, img_path: str = None): + self.calls.append((system, user)) + return DESCRIBE_STUB_YAML, "stop" + + +def _extract_diff_section(user_prompt: str) -> str: + """Return the rendered diff body between the prompt's fixed markers.""" + start_marker = "The PR Git Diff:" + end_marker = "Note that lines in the diff body are prefixed" + start = user_prompt.find(start_marker) + end = user_prompt.find(end_marker) + assert start != -1 and end != -1 and end > start, "describe prompt is missing the diff section" + return user_prompt[start + len(start_marker):end].strip("=\n ").strip() + + +async def _run_smoke() -> None: + pr_url = os.getenv('TEST_PR_URL', DEFAULT_PR_URL) + + get_settings().set("config.git_provider", "github") + get_settings().set("config.publish_output", False) + get_settings().set("config.fallback_models", []) + get_settings().set("config.propagate_tool_errors", True) + + stub = StubDescribeHandler() + agent = PRAgent(ai_handler=lambda: stub) + ok = await agent.handle_request(pr_url, ['describe']) + + assert ok, "describe command failed against the target GitHub PR" + assert stub.calls, "the AI handler was never invoked - describe did not reach the model boundary" + + system_prompt, user_prompt = stub.calls[0] + assert isinstance(system_prompt, str) and system_prompt.strip(), "describe system prompt rendered empty" + diff_section = _extract_diff_section(user_prompt) + assert diff_section, "describe prompt diff section is empty" + + if pr_url == DEFAULT_PR_URL: + assert "azuredevops_provider.py" in diff_section, \ + "describe prompt is missing the default target PR diff content" + assert "Related Ticket Info:" in user_prompt, "describe prompt is missing the related-ticket section" + assert "Ticket Title: 'Azure DevOps get_latest_commit_url" in user_prompt, \ + "expanded ticket data did not reach the describe prompt" + + artifact = dict(get_settings().data).get('artifact', '') + assert isinstance(artifact, str) and artifact.startswith("###"), \ + f"Expected markdown artifact, got {type(artifact).__name__}" + assert "PR Type" in artifact, "PR description artifact missing the PR Type section" + assert "Description" in artifact, "PR description artifact missing the Description section" + + +def test_github_describe_smoke(): + global_settings.get("config") # force the lazy load before copying settings + with request_cycle_context({}): + context['settings'] = copy.deepcopy(global_settings) + asyncio.run(_run_smoke()) + + +if __name__ == '__main__': + test_github_describe_smoke() diff --git a/tests/unittest/test_apply_repo_settings.py b/tests/unittest/test_apply_repo_settings.py index ad8ab37d83..2b7dc50259 100644 --- a/tests/unittest/test_apply_repo_settings.py +++ b/tests/unittest/test_apply_repo_settings.py @@ -92,6 +92,49 @@ def test_repo_without_toml_does_not_inherit_previous_repo_settings( assert "MARKER-FROM-REPO-A" not in _extra_instructions("pr_code_suggestions"), \ "repo A's [pr_code_suggestions].extra_instructions leaked into repo B" + def test_repo_settings_cannot_raise_sibling_context_cap(self, fresh_global_settings, monkeypatch): + """A repo's .pr_agent.toml must not be able to raise + repo_context_max_sibling_files: it bounds cross-repo fetches per build, + so it is host-only. Repo file selection (repo_context_files) may still + be set from the repo, only the cap is refused. + """ + repo_toml = b""" +[config] +repo_context_files = [{repo_id = "group/A/idea", file_path = "AGENTS.md"}] +repo_context_max_sibling_files = 1000 +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: FakeGitProvider(repo_toml), + ) + with request_cycle_context({}): + context["settings"] = copy.deepcopy(global_settings) + git_utils.apply_repo_settings("https://git.example/projects/A/repos/a/pull-requests/1") + assert get_settings().config.get("repo_context_max_sibling_files") == 5, \ + "Repo settings must not be able to raise repo_context_max_sibling_files" + assert get_settings().config.get("repo_context_files") == [ + {"repo_id": "group/A/idea", "file_path": "AGENTS.md"}, + ], "Repo settings should still be able to select repo_context_files" + + def test_repo_selects_siblings_without_changing_host_allowlist(self, fresh_global_settings, monkeypatch): + repo_toml = b""" +[config] +repo_context_sibling_repos = ["group/private"] +repo_context_files = [{repo_id = "group/approved", file_path = "api.py"}] +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: FakeGitProvider(repo_toml), + ) + with request_cycle_context({}): + context["settings"] = copy.deepcopy(global_settings) + get_settings().set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", ["group/approved"]) + git_utils.apply_repo_settings("https://git.example/projects/A/repos/a/pull-requests/1") + assert get_settings().config.repo_context_sibling_repos == ["group/approved"] + assert get_settings().config.repo_context_files == [ + {"repo_id": "group/approved", "file_path": "api.py"}, + ] + def test_unknown_section_does_not_leak_to_next_repo(self, fresh_global_settings, monkeypatch): """Catches the case where a repo's `.pr_agent.toml` introduces a section name not present in the startup defaults. With the per-request clone, diff --git a/tests/unittest/test_azure_devops_provider.py b/tests/unittest/test_azure_devops_provider.py index c856f42785..e1cbc6563c 100644 --- a/tests/unittest/test_azure_devops_provider.py +++ b/tests/unittest/test_azure_devops_provider.py @@ -41,6 +41,109 @@ def test_azure_resolve_comment_thread_closes_thread(): provider.set_thread_status.assert_called_once_with(42, "closed") +def _language_provider(items): + provider = AzureDevopsProvider.__new__(AzureDevopsProvider) + provider.workspace_slug = "proj" + provider.repo_slug = "repo" + provider.pr = SimpleNamespace( + last_merge_target_commit=SimpleNamespace(commit_id="base-sha"), + last_merge_commit=SimpleNamespace(commit_id="head-sha"), + ) + provider.azure_devops_client = MagicMock() + provider.azure_devops_client.get_items.return_value = items + return provider + + +def test_get_languages_returns_language_names(): + # get_languages() must key on language NAMES (e.g. "Python"), not raw + # extensions ("py"): sort_files_by_main_languages() maps names back to + # extensions, so extension keys would drop every file into "Other". + provider = _language_provider([ + SimpleNamespace(git_object_type="blob", path="a.py"), + SimpleNamespace(git_object_type="blob", path="b.py"), + SimpleNamespace(git_object_type="blob", path="c.py"), + SimpleNamespace(git_object_type="blob", path="d.js"), + SimpleNamespace(git_object_type="blob", path="weird.zzz"), + ]) + + languages = provider.get_languages() + + # 3 Python + 1 JavaScript known; .zzz is unknown and excluded from the total. + assert languages == {"Python": 75.0, "JavaScript": 25.0} + + +def test_get_languages_queries_target_commit_inventory(): + provider = _language_provider( + [SimpleNamespace(git_object_type="blob", path="a.py")] + ) + + provider.get_languages() + + call_kwargs = provider.azure_devops_client.get_items.call_args.kwargs + assert call_kwargs["version_descriptor"].version == "base-sha" + assert call_kwargs["version_descriptor"].version_type == "commit" + + +def test_get_languages_returns_empty_map_when_target_commit_missing(): + provider = _language_provider( + [SimpleNamespace(git_object_type="blob", path="a.py")] + ) + provider.pr = SimpleNamespace(last_merge_target_commit=None) + + assert provider.get_languages() == {} + provider.azure_devops_client.get_items.assert_not_called() + + +def test_get_languages_returns_empty_map_when_target_commit_id_missing(): + provider = _language_provider( + [SimpleNamespace(git_object_type="blob", path="a.py")] + ) + provider.pr = SimpleNamespace(last_merge_target_commit=SimpleNamespace(commit_id=None)) + + assert provider.get_languages() == {} + provider.azure_devops_client.get_items.assert_not_called() + + +def test_get_languages_returns_empty_map_when_nothing_matches(): + provider = _language_provider([ + SimpleNamespace(git_object_type="blob", path="weird.zzz"), + SimpleNamespace(git_object_type="blob", path=""), + ]) + + assert provider.get_languages() == {} + + +def test_get_languages_maps_full_paths_and_multipart_extensions(): + # The matcher must classify by the basename and honor multipart extensions. + provider = _language_provider([ + SimpleNamespace(git_object_type="blob", path="src/foo.py"), + SimpleNamespace(git_object_type="blob", path="lib/bar.py"), + SimpleNamespace(git_object_type="blob", path="doc/README.md"), + SimpleNamespace(git_object_type="blob", path="notes.txt"), + SimpleNamespace(git_object_type="blob", path="tpl/file.test.ts"), + ]) + + languages = provider.get_languages() + + assert languages == {"Python": 40.0, "Markdown": 20.0, "Text": 20.0, "TypeScript": 20.0} + + +def test_get_languages_ignores_non_blob_items_and_other_languages(): + # Percentages come from the repository blob inventory, not the PR change set: + # non-blob entries (e.g. folders) must be skipped, and off-language files + # must not skew the ranking. + provider = _language_provider([ + SimpleNamespace(git_object_type="Folder", path="src"), + SimpleNamespace(git_object_type="blob", path="src/app.py"), + SimpleNamespace(git_object_type="blob", path="src/main.py"), + SimpleNamespace(git_object_type="blob", path="render.bin"), + ]) + + languages = provider.get_languages() + + assert languages == {"Python": 100.0} + + class TestAzureDevopsProviderRepoContext: def test_get_repo_file_content_reads_from_target_commit(self): # Repo-context files must be read from the PR target (base) commit, matching diff --git a/tests/unittest/test_bitbucket_provider.py b/tests/unittest/test_bitbucket_provider.py index 229749092d..3cc5562abb 100644 --- a/tests/unittest/test_bitbucket_provider.py +++ b/tests/unittest/test_bitbucket_provider.py @@ -651,6 +651,53 @@ def test_get_issue_comments_normalizes_cloud_comments(self): assert comments[0].body == "## PR Review\n\nreview body" assert comments[0]._cloud_comment is comment + def test_persistent_review_update_links_status_to_existing_cloud_comment(self): + provider = BitbucketProvider.__new__(BitbucketProvider) + provider.pr = MagicMock() + provider.max_comment_length = 32_768 + provider.get_latest_commit_url = MagicMock(return_value="https://bitbucket.org/acme/repo/commits/head") + provider.publish_comment = MagicMock(return_value={"id": 8}) + + header = "## PR Review" + identity = PRReviewIdentity.REGULAR.value + comment_url = "https://bitbucket.org/acme/repo/pull-requests/1#comment-7" + existing = MagicMock() + existing.raw = f"{header}\n\n{identity}\n\nprevious review" + existing.data = {"links": {"html": {"href": comment_url}}} + provider.pr.comments.return_value = [existing] + + provider.publish_persistent_comment( + f"{header}\n\n{identity}\n\nupdated review", + initial_header=header, + final_update_message=True, + identity_marker=identity, + ) + + provider.publish_comment.assert_called_once_with( + f"**[Persistent review]({comment_url})** updated to latest commit " + "https://bitbucket.org/acme/repo/commits/head" + ) + + + def test_get_languages_returns_empty_map_when_repo_has_no_language(self): + # Bitbucket Cloud reports "language": null for repos it cannot classify; + # get_languages() must not emit a {None: 0} key or sort_files_by_main_languages() + # crashes on None.lower() (issue #3340). + repo = MagicMock() + repo.get_data.side_effect = lambda key: None + provider = BitbucketProvider.__new__(BitbucketProvider) + provider.repo = repo + + assert provider.get_languages() == {} + + def test_get_languages_includes_detected_language(self): + repo = MagicMock() + repo.get_data.side_effect = lambda key: "Python" + provider = BitbucketProvider.__new__(BitbucketProvider) + provider.repo = repo + + assert provider.get_languages() == {"Python": 0} + class TestBitbucketServerProvider: @staticmethod def _code_suggestion(line: int): @@ -1089,6 +1136,24 @@ def test_get_repo_file_content_reads_from_target_ref(self): assert content == "repo context" provider.get_file.assert_called_once_with("AGENTS.md", "base-sha") + def test_remove_comment_deletes_comment_with_id_and_version(self): + provider = self._persistent_provider() + comment = SimpleNamespace(id=7, version=3) + + assert provider.remove_comment(comment) is True + + provider.bitbucket_client.delete_pull_request_comment.assert_called_once_with( + "AAA", "my-repo", 1, 7, 3 + ) + + def test_remove_comment_rejects_missing_version(self): + provider = self._persistent_provider() + comment = {"id": 7} + + assert provider.remove_comment(comment) is False + + provider.bitbucket_client.delete_pull_request_comment.assert_not_called() + def test_get_repo_file_content_from_default_branch(self): provider = BitbucketServerProvider.__new__(BitbucketServerProvider) provider.workspace_slug = "AAA" @@ -1130,6 +1195,24 @@ def test_get_repo_file_content_propagates_non_404_errors(self): with pytest.raises(HTTPError, match="500 Internal Server Error"): provider.get_repo_file_content("AGENTS.md") + def test_get_languages_returns_language_names(self): + # get_languages() must key on language NAMES (e.g. "Python"), not raw + # extensions ("py"): sort_files_by_main_languages() maps names back to + # extensions, so extension keys would drop every file into "Other". + provider = BitbucketServerProvider.__new__(BitbucketServerProvider) + provider.get_files = MagicMock(return_value=["a.py", "b.py", "c.py", "d.js", "weird.zzz"]) + + languages = provider.get_languages() + + # 3 Python + 1 JavaScript known; .zzz is unknown and excluded from the total. + assert languages == {"Python": 75.0, "JavaScript": 25.0} + + def test_get_languages_returns_empty_map_when_nothing_matches(self): + provider = BitbucketServerProvider.__new__(BitbucketServerProvider) + provider.get_files = MagicMock(return_value=["weird.zzz", ""]) + + assert provider.get_languages() == {} + def _make_provider_for_repo_settings(self, get_content_side_effect): # Bypass __init__ (which performs live API calls) and only wire up the # attributes get_repo_settings() relies on. diff --git a/tests/unittest/test_certifi_pin_policy.py b/tests/unittest/test_certifi_pin_policy.py new file mode 100644 index 0000000000..2e12c88e0c --- /dev/null +++ b/tests/unittest/test_certifi_pin_policy.py @@ -0,0 +1,34 @@ +"""Guard for the security-pin policy declared above ``dependencies`` in pyproject.toml. + +certifi is the CA trust store every outbound HTTPS call in pr-agent ultimately +validates against. An ``==`` pin freezes that bundle for pip consumers, who then +cannot pick up a rebuilt bundle -- one that drops a distrusted root, say -- without +overriding our requirement. The policy comment in pyproject.toml says +security-sensitive packages declare ranges for exactly this reason (#2523, #3194); +this test is what keeps the comment honest, since nothing else checks it. + +This asserts the declared requirement only. Exact versions for this repo's own +builds still come from uv.lock, so reproducibility is unaffected. +""" +import tomllib + +CERTIFI_POLICY = ( + "certifi must declare a floor (>=), not an == pin: it is the CA trust store, and an exact " + "pin leaves pip consumers stuck on a stale certificate bundle. Exact versions for this " + "repo's own builds come from uv.lock. See the comment above [project].dependencies." +) + + +def _certifi_requirement(pytestconfig): + with open(pytestconfig.inipath, "rb") as f: + dependencies = tomllib.load(f)["project"]["dependencies"] + + matches = [dep for dep in dependencies if dep.lower().startswith("certifi")] + assert len(matches) == 1, f"expected exactly one certifi requirement, found {matches!r}" + return matches[0] + + +def test_certifi_is_not_exact_pinned(pytestconfig): + requirement = _certifi_requirement(pytestconfig) + assert "==" not in requirement, f"{requirement!r}: {CERTIFI_POLICY}" + assert ">=" in requirement, f"{requirement!r}: {CERTIFI_POLICY}" diff --git a/tests/unittest/test_cli_args_security.py b/tests/unittest/test_cli_args_security.py index 1f775d4bd9..77bf752102 100644 --- a/tests/unittest/test_cli_args_security.py +++ b/tests/unittest/test_cli_args_security.py @@ -68,6 +68,21 @@ "--pr_reviewer.publish_error_details=true", "--pr_reviewer__publish_error_details=true", '--pr_reviewer={"publish_error_details": true}', + # repo_context_max_sibling_files is host-only: letting a comment raise it would defeat the + # sibling-fetch safety bound and allow unbounded cross-repository API calls. + "--config.repo_context_max_sibling_files=1000", + "--config__repo_context_max_sibling_files=1000", + '--config={"repo_context_max_sibling_files": 1000}', + # repo_context_files selects which files become model instructions, so comment arguments + # must not be able to point the bot at arbitrary sibling repo content. + '--config.repo_context_files=[{"repo_id": "group/A/idea", "file_path": "AGENTS.md"}]', + "--config__repo_context_files=[\"AGENTS.md\"]", + '--config={"repo_context_files": ["AGENTS.md"]}', + # repo_context_sibling_repos is the host-only allowlist of sibling repositories whose + # files may be selected; neither repo settings nor comment arguments can change it. + "--config.repo_context_sibling_repos=[]", + "--config__repo_context_sibling_repos=[]", + '--config={"repo_context_sibling_repos": []}', ] diff --git a/tests/unittest/test_cli_artifact_context.py b/tests/unittest/test_cli_artifact_context.py index 982807a076..50d6322bb0 100644 --- a/tests/unittest/test_cli_artifact_context.py +++ b/tests/unittest/test_cli_artifact_context.py @@ -7,7 +7,7 @@ def test_run_injects_the_artifact_context_before_handling_the_request(): """A pipeline that runs the CLI gets the same [artifacts] injection as the GitHub Action.""" order = [] - fake_settings = SimpleNamespace(litellm={}, set=MagicMock()) + fake_settings = SimpleNamespace(config={}, litellm={}, set=MagicMock()) async def fake_handle_request(*_args, **_kwargs): order.append("handle_request") diff --git a/tests/unittest/test_cli_config_branch.py b/tests/unittest/test_cli_config_branch.py index c3f0e09b69..e04a511433 100644 --- a/tests/unittest/test_cli_config_branch.py +++ b/tests/unittest/test_cli_config_branch.py @@ -11,6 +11,7 @@ def test_set_parser_supports_config_branch_flag(): def test_run_sets_config_branch_from_cli_flag(): fake_settings = SimpleNamespace( + config={}, litellm={}, set=MagicMock(), ) @@ -29,6 +30,7 @@ async def fake_handle_request(*_args, **_kwargs): def test_run_sets_config_branch_from_env_var(): fake_settings = SimpleNamespace( + config={}, litellm={}, set=MagicMock(), ) @@ -51,6 +53,7 @@ async def fake_handle_request(*_args, **_kwargs): def test_run_whitespace_cli_branch_falls_back_to_env_var(): """A whitespace-only --config-branch must not short-circuit the env fallback.""" fake_settings = SimpleNamespace( + config={}, litellm={}, set=MagicMock(), ) @@ -74,6 +77,7 @@ def test_run_reconciles_config_branch_when_absent(): """Without a flag or env var, CONFIG.CONFIG_BRANCH must be reset (not leaked from a previous run() call in the same process-wide settings singleton).""" fake_settings = SimpleNamespace( + config={}, litellm={}, set=MagicMock(), ) diff --git a/tests/unittest/test_cli_exit_status.py b/tests/unittest/test_cli_exit_status.py new file mode 100644 index 0000000000..199e0e762e --- /dev/null +++ b/tests/unittest/test_cli_exit_status.py @@ -0,0 +1,431 @@ +import os +import subprocess +import sys +import tomllib +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from pr_agent import cli +from pr_agent.config_loader import get_settings + +_CONSOLE_SCRIPT_FALLBACK = ( + "import sys; from pr_agent.cli import run; sys.exit(run())" +) +_SETTINGS_KEYS = [ + "config.propagate_tool_errors", + "config.cli_mode", + "config.config_branch", + "config.extra_config_url", +] + + +@pytest.fixture(autouse=True) +def restore_cli_settings(): + settings = get_settings() + original = {key: settings.get(key, None) for key in _SETTINGS_KEYS} + yield + for key, value in original.items(): + settings.set(key, value) + + +def _run_with_result(monkeypatch, result, *, propagate_tool_errors): + fake_settings = SimpleNamespace( + config={"propagate_tool_errors": propagate_tool_errors}, + litellm={}, + set=MagicMock(), + ) + + async def fake_handle_request(*_args, **_kwargs): + return result + + monkeypatch.setattr(cli, "get_settings", lambda: fake_settings) + monkeypatch.setattr( + cli, + "PRAgent", + lambda: SimpleNamespace(handle_request=fake_handle_request), + ) + monkeypatch.setattr(cli, "inject_artifact_context", lambda: None) + + return cli.run(inargs=["--pr_url=https://example.com/org/repo/pull/1", "review"]) + + +@pytest.mark.parametrize( + ("request_result", "propagate_tool_errors", "expected_status", "prints_help"), + [ + (True, False, None, False), + (True, True, None, False), + (False, False, None, True), + (False, True, 1, True), + (None, True, None, True), + (0, True, None, True), + ], +) +def test_run_maps_request_result_to_status( + monkeypatch, + capsys, + request_result, + propagate_tool_errors, + expected_status, + prints_help, +): + status = _run_with_result( + monkeypatch, + request_result, + propagate_tool_errors=propagate_tool_errors, + ) + + assert status == expected_status + assert ("usage:" in capsys.readouterr().out) is prints_help + + +def test_run_reads_effective_setting_after_dispatch(monkeypatch): + fake_settings = SimpleNamespace( + config={"propagate_tool_errors": False}, + litellm={}, + set=MagicMock(), + ) + + async def fake_handle_request(*_args, **_kwargs): + fake_settings.config["propagate_tool_errors"] = True + return False + + monkeypatch.setattr(cli, "get_settings", lambda: fake_settings) + monkeypatch.setattr( + cli, + "PRAgent", + lambda: SimpleNamespace(handle_request=fake_handle_request), + ) + monkeypatch.setattr(cli, "inject_artifact_context", lambda: None) + + assert cli.run(inargs=["--pr_url=https://example.com/org/repo/pull/1", "review"]) == 1 + + +def test_run_drains_callbacks_before_returning_failure_status(monkeypatch): + events = [] + fake_settings = SimpleNamespace( + config={"propagate_tool_errors": True}, + litellm={}, + set=MagicMock(), + ) + + async def fake_handle_request(*_args, **_kwargs): + events.append("request") + return False + + async def fake_drain(*_args, **_kwargs): + events.append("drain") + + monkeypatch.setattr(cli, "get_settings", lambda: fake_settings) + monkeypatch.setattr( + cli, + "PRAgent", + lambda: SimpleNamespace(handle_request=fake_handle_request), + ) + monkeypatch.setattr(cli, "inject_artifact_context", lambda: None) + monkeypatch.setattr(cli, "litellm_callbacks_registered", lambda: True) + monkeypatch.setattr(cli, "drain_litellm_callbacks", fake_drain) + + status = cli.run(inargs=["--pr_url=https://example.com/org/repo/pull/1", "review"]) + + events.append(f"status:{status}") + assert events == ["request", "drain", "status:1"] + + +@pytest.mark.parametrize( + ("repo_value", "cli_value", "expected_status"), + [ + (False, True, 1), + (True, None, 1), + (True, False, None), + ], +) +def test_real_request_failure_uses_effective_propagation_setting( + monkeypatch, + repo_value, + cli_value, + expected_status, +): + from pr_agent.agent import pr_agent as pr_agent_module + + events = [] + expected_effective_value = repo_value if cli_value is None else cli_value + + class FailingReview: + def __init__(self, *_args, **_kwargs): + pass + + async def run(self): + events.append("tool") + assert ( + get_settings().config.get("propagate_tool_errors") + is expected_effective_value + ) + raise RuntimeError("controlled tool failure") + + async def fake_drain(*_args, **_kwargs): + events.append("callbacks") + + def fake_apply_repo_settings(*_args, **_kwargs): + get_settings().set("CONFIG.PROPAGATE_TOOL_ERRORS", repo_value) + + monkeypatch.setitem(pr_agent_module.command2class, "review", FailingReview) + monkeypatch.setattr(pr_agent_module, "apply_repo_settings", fake_apply_repo_settings) + monkeypatch.setattr(pr_agent_module, "flush_telemetry", lambda: events.append("telemetry")) + monkeypatch.setattr(cli, "inject_artifact_context", lambda: None) + monkeypatch.setattr(cli, "litellm_callbacks_registered", lambda: True) + monkeypatch.setattr(cli, "drain_litellm_callbacks", fake_drain) + + inargs = [ + "--pr_url=https://example.com/org/repo/pull/1", + "review", + ] + if cli_value is not None: + inargs.append(f"--config.propagate_tool_errors={str(cli_value).lower()}") + + status = cli.run(inargs=inargs) + + events.append(f"status:{status}") + assert events == [ + "tool", + "telemetry", + "callbacks", + f"status:{expected_status}", + ] + + +@pytest.mark.parametrize( + ("baseline", "override_source", "override", "first_succeeds", "expected_statuses"), + [ + (False, "repository", True, False, [1, None]), + (False, "command", True, False, [1, None]), + (True, "repository", False, False, [None, 1]), + (True, "command", False, False, [None, 1]), + (False, "repository", True, True, [None, None]), + (False, "command", True, True, [None, None]), + (True, "repository", False, True, [None, 1]), + (True, "command", False, True, [None, 1]), + ], +) +def test_run_restores_propagation_setting_between_invocations( + monkeypatch, + baseline, + override_source, + override, + first_succeeds, + expected_statuses, +): + from pr_agent.agent import pr_agent as pr_agent_module + + settings = get_settings() + settings.set("CONFIG.PROPAGATE_TOOL_ERRORS", baseline) + apply_calls = 0 + observed_values = [] + request_results = iter([first_succeeds, False]) + + class ControlledReview: + def __init__(self, *_args, **_kwargs): + pass + + async def run(self): + observed_values.append(settings.config.get("propagate_tool_errors")) + if next(request_results): + return + raise RuntimeError("controlled tool failure") + + def fake_apply_repo_settings(*_args, **_kwargs): + nonlocal apply_calls + if override_source == "repository" and apply_calls == 0: + settings.set("CONFIG.PROPAGATE_TOOL_ERRORS", override) + apply_calls += 1 + + monkeypatch.setitem(pr_agent_module.command2class, "review", ControlledReview) + monkeypatch.setattr(pr_agent_module, "apply_repo_settings", fake_apply_repo_settings) + monkeypatch.setattr(pr_agent_module, "flush_telemetry", lambda: None) + monkeypatch.setattr(cli, "inject_artifact_context", lambda: None) + monkeypatch.setattr(cli, "litellm_callbacks_registered", lambda: False) + + first_inargs = ["--pr_url=https://example.com/org/repo/pull/1", "review"] + if override_source == "command": + first_inargs.append(f"--config.propagate_tool_errors={str(override).lower()}") + + statuses = [ + cli.run(inargs=first_inargs), + cli.run(inargs=["--pr_url=https://example.com/org/repo/pull/2", "review"]), + ] + + assert observed_values == [override, baseline] + assert statuses == expected_statuses + assert settings.config.get("propagate_tool_errors") is baseline + + +def test_run_restores_propagation_setting_when_callback_drain_raises(monkeypatch): + settings = get_settings() + settings.set("CONFIG.PROPAGATE_TOOL_ERRORS", False) + + async def fake_handle_request(*_args, **_kwargs): + settings.set("CONFIG.PROPAGATE_TOOL_ERRORS", True) + return False + + async def failing_drain(*_args, **_kwargs): + raise RuntimeError("controlled callback failure") + + monkeypatch.setattr( + cli, + "PRAgent", + lambda: SimpleNamespace(handle_request=fake_handle_request), + ) + monkeypatch.setattr(cli, "inject_artifact_context", lambda: None) + monkeypatch.setattr(cli, "litellm_callbacks_registered", lambda: True) + monkeypatch.setattr(cli, "drain_litellm_callbacks", failing_drain) + + with pytest.raises(RuntimeError, match="controlled callback failure"): + cli.run(inargs=["--pr_url=https://example.com/org/repo/pull/1", "review"]) + + assert settings.config.get("propagate_tool_errors") is False + + +def _console_script_entrypoint(python_executable): + console_script = Path(python_executable).with_name("pr-agent") + if sys.platform == "win32": + console_script = console_script.with_suffix(".exe") + if console_script.is_file(): + return [str(console_script)] + return [python_executable, "-c", _CONSOLE_SCRIPT_FALLBACK] + + +def test_console_script_maps_to_cli_run(): + with (Path(__file__).parents[2] / "pyproject.toml").open("rb") as pyproject: + project = tomllib.load(pyproject)["project"] + + assert project["scripts"]["pr-agent"] == "pr_agent.cli:run" + + +def test_console_script_falls_back_in_source_only_environment(tmp_path): + python_executable = str(tmp_path / "bin" / "python") + assert _console_script_entrypoint(python_executable) == [ + python_executable, + "-c", + _CONSOLE_SCRIPT_FALLBACK, + ] + + +def test_run_command_forwards_run_status(monkeypatch): + def fake_run(*, args): + assert args.pr_url == "https://example.com/org/repo/pull/1" + assert args.command == "review" + return 1 + + monkeypatch.setattr(cli, "run", fake_run) + + assert cli.run_command("https://example.com/org/repo/pull/1", "/review") == 1 + + +@pytest.fixture +def process_entrypoints(): + return ( + _console_script_entrypoint(sys.executable), + [sys.executable, "-m", "pr_agent.cli"], + [sys.executable, "-m", "pr_agent.cli_pip"], + [sys.executable, "pr_agent/cli_pip.py"], + ) + + +@pytest.fixture +def argparse_entrypoints(process_entrypoints): + return process_entrypoints[:2] + + +@pytest.fixture +def controlled_process_env(tmp_path): + sitecustomize = tmp_path / "sitecustomize.py" + sitecustomize.write_text( + """ +import os + +from pr_agent.agent.pr_agent import PRAgent +from pr_agent.config_loader import get_settings + + +async def controlled_handle_request(self, pr_url, request, notify=None): + get_settings().set( + \"CONFIG.PROPAGATE_TOOL_ERRORS\", + os.environ.get(\"PR_AGENT_TEST_PROPAGATE\") == \"true\", + ) + result = os.environ.get(\"PR_AGENT_TEST_RESULT\") + if result == \"none\": + return None + if result == \"zero\": + return 0 + return result == \"true\" + + +PRAgent.handle_request = controlled_handle_request +""".lstrip(), + encoding="utf-8", + ) + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + part for part in (str(tmp_path), env.get("PYTHONPATH")) if part + ) + return env + + +@pytest.mark.parametrize( + ("request_result", "propagate_tool_errors", "expected_status"), + [ + ("false", "true", 1), + ("false", "false", 0), + ], +) +def test_process_entrypoints_map_request_status( + process_entrypoints, + controlled_process_env, + request_result, + propagate_tool_errors, + expected_status, +): + env = controlled_process_env.copy() + env["PR_AGENT_TEST_RESULT"] = request_result + env["PR_AGENT_TEST_PROPAGATE"] = propagate_tool_errors + + for entrypoint in process_entrypoints: + completed = subprocess.run( + [*entrypoint, "--pr_url=https://example.com/org/repo/pull/1", "review"], + cwd=Path(__file__).parents[2], + env=env, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert completed.returncode == expected_status, ( + entrypoint, + completed.stdout, + completed.stderr, + ) + + +def test_process_entrypoints_preserve_argparse_status( + argparse_entrypoints, + controlled_process_env, +): + for entrypoint in argparse_entrypoints: + completed = subprocess.run( + [*entrypoint, "not-a-command"], + cwd=Path(__file__).parents[2], + env=controlled_process_env, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert completed.returncode == 2, ( + entrypoint, + completed.stdout, + completed.stderr, + ) diff --git a/tests/unittest/test_cli_plain_diff_mode.py b/tests/unittest/test_cli_plain_diff_mode.py index 2011d1f1d8..b3a9a057e4 100644 --- a/tests/unittest/test_cli_plain_diff_mode.py +++ b/tests/unittest/test_cli_plain_diff_mode.py @@ -1,13 +1,119 @@ +import builtins +import io + import pytest -from pr_agent.cli import run, set_parser +from pr_agent.cli import _resolve_output_option, commands, run, set_parser from pr_agent.config_loader import get_settings # Keys run() mutates on the process-wide settings singleton, directly or via the # diff-mode CLI path. Snapshotted and restored around every test (autouse) so # state never leaks, even when run() sets keys the test never touches itself. -_SETTINGS_KEYS = ["plain_diff.content", "plain_diff.output_path", "plain_diff.json_output_path", - "config.git_provider", "config.publish_output"] +_SETTINGS_KEYS = [ + "plain_diff.content", + "plain_diff.output_path", + "plain_diff.json_output_path", + "config.git_provider", + "config.publish_output", + "config.cli_mode", + "config.config_branch", + "config.extra_config_url", +] + +_DIFF = ( + "diff --git a/foo.py b/foo.py\n" + "index 1111111..2222222 100644\n" + "--- a/foo.py\n" + "+++ b/foo.py\n" + "@@ -1,3 +1,3 @@\n" + " line1\n-line2\n+line2-changed\n line3\n" +) +_REVIEW_MARKER = "plain-diff auto-review output" +_CANNED_REVIEW = f"""\ +review: + estimated_effort_to_review_[1-5]: '1' + score: '90' + relevant_tests: 'No' + key_issues_to_review: + - relevant_file: foo.py + issue_header: '{_REVIEW_MARKER}' + issue_content: 'Verify the changed line.' + start_line: 2 + end_line: 2 + security_concerns: 'No' +""" +_OUTPUT_OPTION_PREFIXES = [ + (option[:length], option) + for option in ("--output", "--json-output") + for length in range(3, len(option) + 1) +] + +_MARKDOWN_COMMANDS = [ + "review", + "review_pr", + "auto_review", + "describe", + "describe_pr", + "improve", + "improve_code", + "ask", + "ask_question", + "config", + "settings", + "help", +] + +_NON_MARKDOWN_COMMANDS = [ + "answer", + "ask_line", + "update_changelog", + "similar_issue", + "add_docs", + "generate_labels", +] + + +def test_output_command_matrix_covers_cli_commands(): + assert set(_MARKDOWN_COMMANDS) | set(_NON_MARKDOWN_COMMANDS) == set(commands) + assert set(_MARKDOWN_COMMANDS).isdisjoint(_NON_MARKDOWN_COMMANDS) + + +def _rejecting_input_args(input_mode, monkeypatch): + class UnreadableStdin: + def read(self): + pytest.fail("stdin must not be read for invalid output options") + + if input_mode == "stdin": + monkeypatch.setattr("sys.stdin", UnreadableStdin()) + return ["--stdin"] + + original_open = builtins.open + + def guarded_open(path, *args, **kwargs): + if path == "changes.diff": + pytest.fail("the diff file must not be opened for invalid output options") + return original_open(path, *args, **kwargs) + + monkeypatch.setattr("builtins.open", guarded_open) + return ["--diff-file", "changes.diff"] + + +def _valid_input_args(input_mode, monkeypatch, tmp_path): + if input_mode == "stdin": + monkeypatch.setattr("sys.stdin", io.StringIO(_DIFF)) + return ["--stdin"] + + diff_file = tmp_path / "changes.diff" + diff_file.write_text(_DIFF, encoding="utf-8") + return ["--diff-file", str(diff_file)] + + +def _fail_if_agent_constructed(monkeypatch): + class NeverAgent: + def __init__(self): + pytest.fail("PRAgent must not be constructed for invalid output options") + + monkeypatch.setattr("pr_agent.cli.PRAgent", NeverAgent) @pytest.fixture(autouse=True) @@ -43,6 +149,28 @@ def test_parser_stdin_flag(): assert args.stdin is True +@pytest.mark.parametrize( + ("option_spelling", "option"), + _OUTPUT_OPTION_PREFIXES, +) +def test_parser_accepts_unambiguous_output_abbreviations_before_command( + option_spelling, option, +): + parser = set_parser() + args = parser.parse_args(["--stdin", option_spelling, "result", "review"]) + + destination = option[2:].replace("-", "_") + assert getattr(args, destination) == "result" + + +def test_misplaced_output_abbreviation_must_be_unambiguous(): + parser = set_parser() + parser.add_argument("--outcome") + + assert _resolve_output_option(parser, "--out=value") is None + assert _resolve_output_option(parser, "--output=value") == "--output" + + def test_missing_diff_file_fails_fast(tmp_path, capsys): """A non-existent --diff-file must exit cleanly via parser.error (SystemExit) with a clear message, not crash with an uncaught OSError traceback.""" @@ -56,27 +184,232 @@ def test_missing_diff_file_fails_fast(tmp_path, capsys): def test_json_output_outside_diff_mode_fails_fast(capsys): """Reject --json-output in hosted-provider mode via parser.error instead of silently dropping the explicitly requested artifact.""" - with pytest.raises(SystemExit): + with pytest.raises(SystemExit) as exc_info: run(inargs=["--pr_url", "https://example/pr/1", "--json-output", "out.json", "review"]) + assert exc_info.value.code == 2 err = capsys.readouterr().err assert "--json-output is only supported in plain-diff mode" in err -_DIFF = ( - "diff --git a/foo.py b/foo.py\n" - "index 1111111..2222222 100644\n" - "--- a/foo.py\n" - "+++ b/foo.py\n" - "@@ -1,3 +1,3 @@\n" - " line1\n-line2\n+line2-changed\n line3\n" +@pytest.mark.parametrize( + "target_args", + [ + ["--pr_url", "https://example/pr/1"], + ["--issue_url", "https://example/issue/1"], + [], + ], ) +def test_markdown_output_outside_plain_diff_fails_before_dispatch(target_args, monkeypatch, capsys): + _fail_if_agent_constructed(monkeypatch) + + with pytest.raises(SystemExit) as exc_info: + run(inargs=[*target_args, "--output", "out.md", "review"]) + + assert exc_info.value.code == 2 + assert "--output is only supported in plain-diff mode" in capsys.readouterr().err + + +def test_markdown_output_rejects_local_git_mode_before_dispatch(monkeypatch, capsys): + _fail_if_agent_constructed(monkeypatch) + + with pytest.raises(SystemExit) as exc_info: + run(inargs=[ + "--pr_url", "main", + "--output", "out.md", + "review", + "--config.git_provider=local", + ]) + + assert exc_info.value.code == 2 + assert "--output is only supported in plain-diff mode" in capsys.readouterr().err + + +@pytest.mark.parametrize("option", ["--output", "--json-output"]) +@pytest.mark.parametrize("spelling", ["equals", "separate"]) +@pytest.mark.parametrize("mode", ["hosted", "stdin"]) +def test_empty_output_paths_fail_before_input_or_dispatch( + option, spelling, mode, monkeypatch, capsys, +): + _fail_if_agent_constructed(monkeypatch) + target_args = ["--pr_url", "https://example/pr/1"] + if mode == "stdin": + target_args = _rejecting_input_args("stdin", monkeypatch) + option_args = [f"{option}="] if spelling == "equals" else [option, ""] + + with pytest.raises(SystemExit) as exc_info: + run(inargs=[*target_args, *option_args, "review"]) + + assert exc_info.value.code == 2 + assert f"{option} requires a non-empty path" in capsys.readouterr().err + + +@pytest.mark.parametrize("command", _NON_MARKDOWN_COMMANDS) +@pytest.mark.parametrize("input_mode", ["stdin", "file"]) +def test_markdown_output_rejects_unsupported_plain_diff_commands_before_read( + command, input_mode, monkeypatch, capsys, +): + input_args = _rejecting_input_args(input_mode, monkeypatch) + + with pytest.raises(SystemExit) as exc_info: + run(inargs=[*input_args, "--output", "out.md", command]) + + assert exc_info.value.code == 2 + assert "--output is not supported for" in capsys.readouterr().err + + +@pytest.mark.parametrize("command", _MARKDOWN_COMMANDS) +@pytest.mark.parametrize("input_mode", ["stdin", "file"]) +def test_markdown_output_accepts_compatible_plain_diff_commands( + command, input_mode, monkeypatch, tmp_path, +): + captured = {} + + class FakeAgent: + async def handle_request(self, target, request, notify=None): + captured["target"] = target + captured["request"] = request + captured["output_path"] = get_settings().plain_diff.output_path + return True + + monkeypatch.setattr("pr_agent.cli.PRAgent", FakeAgent) + input_args = _valid_input_args(input_mode, monkeypatch, tmp_path) + + output = tmp_path / "result.md" + run(inargs=[*input_args, "--output", str(output), command]) + + assert captured == { + "target": "local_diff", + "request": [command], + "output_path": str(output), + } + + +@pytest.mark.parametrize("command", ["config", "settings", "help"]) +def test_compatible_non_model_commands_write_markdown_artifacts( + command, monkeypatch, tmp_path, +): + monkeypatch.setattr("sys.stdin", io.StringIO(_DIFF)) + output = tmp_path / f"{command}.md" + + run(inargs=["--stdin", "--output", str(output), command]) + + assert output.read_text(encoding="utf-8").strip() + + +def test_auto_review_writes_markdown_artifact_with_stubbed_model( + monkeypatch, tmp_path, +): + async def fake_chat_completion(self, model, system, user, temperature=0.2, **kwargs): + return _CANNED_REVIEW, "stop" + + monkeypatch.setattr( + "pr_agent.algo.ai_handlers.litellm_ai_handler.LiteLLMAIHandler.chat_completion", + fake_chat_completion, + ) + monkeypatch.setattr("sys.stdin", io.StringIO(_DIFF)) + output = tmp_path / "auto-review.md" + + run(inargs=["--stdin", "--output", str(output), "auto_review"]) + + assert _REVIEW_MARKER in output.read_text(encoding="utf-8") + + +def test_answer_output_is_rejected_without_creating_artifact(monkeypatch, tmp_path, capsys): + input_args = _rejecting_input_args("stdin", monkeypatch) + output = tmp_path / "answer.md" + + with pytest.raises(SystemExit) as exc_info: + run(inargs=[*input_args, "--output", str(output), "answer"]) + + assert exc_info.value.code == 2 + assert "--output is not supported for" in capsys.readouterr().err + assert not output.exists() + + +@pytest.mark.parametrize("command", ["review", "review_pr"]) +@pytest.mark.parametrize("input_mode", ["stdin", "file"]) +def test_json_output_accepts_review_aliases(command, input_mode, monkeypatch, tmp_path): + captured = {} + + class FakeAgent: + async def handle_request(self, target, request, notify=None): + captured["request"] = request + captured["json_output_path"] = get_settings().plain_diff.json_output_path + return True + + monkeypatch.setattr("pr_agent.cli.PRAgent", FakeAgent) + input_args = _valid_input_args(input_mode, monkeypatch, tmp_path) + + output = tmp_path / "review.json" + run(inargs=[*input_args, "--json-output", str(output), command]) + + assert captured == { + "request": [command], + "json_output_path": str(output), + } + + +@pytest.mark.parametrize("command", [*_NON_MARKDOWN_COMMANDS, *_MARKDOWN_COMMANDS[2:]]) +@pytest.mark.parametrize("input_mode", ["stdin", "file"]) +def test_json_output_rejects_non_review_commands_before_read( + command, input_mode, monkeypatch, capsys, +): + input_args = _rejecting_input_args(input_mode, monkeypatch) + + with pytest.raises(SystemExit) as exc_info: + run(inargs=[*input_args, "--json-output", "out.json", command]) + + assert exc_info.value.code == 2 + assert "--json-output is only supported for plain-diff review commands" in capsys.readouterr().err + + +@pytest.mark.parametrize( + ("option_spelling", "option"), + _OUTPUT_OPTION_PREFIXES, +) +@pytest.mark.parametrize("value_form", ["separate", "equals"]) +def test_output_options_after_command_fail_before_read( + option_spelling, option, value_form, monkeypatch, capsys, +): + class UnreadableStdin: + def read(self): + pytest.fail("stdin must not be read for misplaced output options") + + monkeypatch.setattr("sys.stdin", UnreadableStdin()) + value = "out.json" if option == "--json-output" else "out.md" + trailing_args = ( + [option_spelling, value] + if value_form == "separate" + else [f"{option_spelling}={value}"] + ) + + with pytest.raises(SystemExit) as exc_info: + run(inargs=["--stdin", "review", *trailing_args]) + + assert exc_info.value.code == 2 + err = capsys.readouterr().err + assert f"{option} must appear before the command" in err + + +def test_near_match_after_command_remains_a_tool_argument(monkeypatch): + captured = {} + + class FakeAgent: + async def handle_request(self, target, request, notify=None): + captured["request"] = request + return True + + monkeypatch.setattr("pr_agent.cli.PRAgent", FakeAgent) + monkeypatch.setattr("sys.stdin", io.StringIO(_DIFF)) + + run(inargs=["--stdin", "review", "--output-format=markdown"]) + + assert captured["request"] == ["review", "--output-format=markdown"] def test_diff_mode_forces_publish_output(cfg, monkeypatch): """Diff mode must force config.publish_output=True so stdout/--output is never suppressed by a config/env that disabled publishing.""" - import io - cfg("config.publish_output", False) captured = {} @@ -92,8 +425,6 @@ async def handle_request(self, target, request, notify=None): def test_diff_mode_sets_json_output_path(cfg, monkeypatch, tmp_path): - import io - captured = {} class FakeAgent: diff --git a/tests/unittest/test_codecommit_client.py b/tests/unittest/test_codecommit_client.py index caa975bd07..3380b96e8d 100644 --- a/tests/unittest/test_codecommit_client.py +++ b/tests/unittest/test_codecommit_client.py @@ -164,3 +164,63 @@ def test_get_pr_preserves_all_target_repositories(self): ("repo-one", "source-one", "destination-one"), ("repo-two", "source-two", "destination-two"), ] + + def test_get_comments_for_pull_request_uses_safe_page_size_and_all_pages(self): + api = CodeCommitClient() + api.boto_client = MagicMock() + paginator = api.boto_client.get_paginator.return_value + paginator.paginate.return_value = [ + {"commentsForPullRequestData": [{"comments": [{"commentId": "comment-1"}]}]}, + {"commentsForPullRequestData": [{"comments": [{"commentId": "comment-2"}]}]}, + ] + + comments = api.get_comments_for_pull_request(321) + + assert comments == [ + {"comments": [{"commentId": "comment-1"}]}, + {"comments": [{"commentId": "comment-2"}]}, + ] + api.boto_client.get_paginator.assert_called_once_with("get_comments_for_pull_request") + paginator.paginate.assert_called_once_with( + pullRequestId="321", + PaginationConfig={"PageSize": 100}, + ) + + def test_publish_comment_returns_boto_response(self): + api = CodeCommitClient() + api.boto_client = MagicMock() + api.boto_client.post_comment_for_pull_request.return_value = { + "comment": {"commentId": "comment-1", "content": "Review"} + } + + response = api.publish_comment( + repo_name="my_test_repo", + pr_number=321, + destination_commit="destination-commit", + source_commit="source-commit", + comment="Review", + ) + + assert response == {"comment": {"commentId": "comment-1", "content": "Review"}} + api.boto_client.post_comment_for_pull_request.assert_called_once_with( + pullRequestId="321", + repositoryName="my_test_repo", + beforeCommitId="destination-commit", + afterCommitId="source-commit", + content="Review", + ) + + def test_update_comment_returns_boto_response(self): + api = CodeCommitClient() + api.boto_client = MagicMock() + api.boto_client.update_comment.return_value = { + "comment": {"commentId": "comment-1", "content": "Updated review"} + } + + response = api.update_comment("comment-1", "Updated review") + + assert response == {"comment": {"commentId": "comment-1", "content": "Updated review"}} + api.boto_client.update_comment.assert_called_once_with( + commentId="comment-1", + content="Updated review", + ) diff --git a/tests/unittest/test_codecommit_provider.py b/tests/unittest/test_codecommit_provider.py index 09a42e5782..117cbe5b93 100644 --- a/tests/unittest/test_codecommit_provider.py +++ b/tests/unittest/test_codecommit_provider.py @@ -1,9 +1,13 @@ +from datetime import datetime, timezone +from types import SimpleNamespace from unittest.mock import MagicMock, call, patch import pytest from pr_agent.algo.types import EDIT_TYPE +from pr_agent.algo.utils import PRReviewHeader, PRReviewIdentity from pr_agent.git_providers.codecommit_provider import CodeCommitFile, CodeCommitProvider, PullRequestCCMimic +from pr_agent.tools.pr_reviewer import PRReviewer class TestCodeCommitFile: @@ -39,6 +43,74 @@ def _make_diff_provider(git_files): provider.codecommit_client = MagicMock() return provider + @staticmethod + def _make_persistent_provider(targets=None): + provider = object.__new__(CodeCommitProvider) + provider.repo_name = "source-repository" + provider.pr_num = 321 + provider.pr_url = ( + "https://us-east-1.console.aws.amazon.com/codesuite/codecommit/" + "repositories/source-repository/pull-requests/321" + ) + provider.diff_files = None + provider.git_files = [] + provider.codecommit_client = MagicMock() + targets = targets or [ + SimpleNamespace( + repository_name="source-repository", + source_commit="source-commit-1", + source_branch="refs/heads/feature", + destination_commit="destination-commit-1", + destination_branch="refs/heads/main", + ) + ] + provider.pr = PullRequestCCMimic("Persistent PR", [], targets=targets) + provider.pr.source_commit = targets[0].source_commit + provider.pr.source_branch = targets[0].source_branch + provider.pr.destination_commit = targets[0].destination_commit + provider.pr.destination_branch = targets[0].destination_branch + return provider + + @staticmethod + def _comment( + body, + comment_id="comment-1", + created_at=None, + last_modified_at=None, + **overrides, + ): + comment = { + "commentId": comment_id, + "content": body, + "creationDate": created_at or datetime(2024, 1, 1, tzinfo=timezone.utc), + "lastModifiedDate": last_modified_at or created_at or datetime(2024, 1, 1, tzinfo=timezone.utc), + "authorArn": "arn:aws:iam::123456789012:user/pr-agent", + } + comment.update(overrides) + return comment + + @staticmethod + def _comment_group( + *comments, + repository_name="source-repository", + before_commit_id="old-destination-commit", + after_commit_id="old-source-commit", + location=None, + ): + group = { + "repositoryName": repository_name, + "beforeCommitId": before_commit_id, + "afterCommitId": after_commit_id, + "comments": list(comments), + } + if location is not None: + group["location"] = location + return group + + @staticmethod + def _successful_update_response(comment_id, body): + return {"comment": {"commentId": comment_id, "content": body, "deleted": False}} + def test_get_diff_files_includes_deleted_file(self): provider = object.__new__(CodeCommitProvider) provider.repo_name = "my_test_repo" @@ -200,6 +272,76 @@ def test_get_files_includes_differences_from_every_pull_request_target(self): ("two.py", "before two\n", "after two\n"), ] + def test_prepare_comment_body_caps_at_codecommit_limit(self): + # PostCommentForPullRequest rejects bodies above 10,240 characters and + # publish_comment raises instead of degrading (#3272). The cap must be + # measured AFTER the newline doubling, which grows the body. + provider = self._make_persistent_provider() + body = "\n".join(["x" * 100] * 120) # 12,099 chars before doubling + + prepared = provider._prepare_comment_body(body) + + assert len(prepared) <= 10240 + assert prepared.endswith("...") + assert "\n\n" in prepared + + def test_prepare_comment_body_keeps_review_state_marker_parseable(self): + # The reviewer budgets the hidden state marker before the newline + # doubling; once doubled, the body can exceed the cap. Only the human + # text may be truncated, or the next run cannot read the state. + from pr_agent.algo.review_finding_state import append_review_state, parse_review_state + + provider = self._make_persistent_provider() + state = { + "schema_version": 1, + "last_run": {"commit": "abc123"}, + "findings": [ + {"finding_id": f"f{i}", "state": "ACTIVE", "path": "a.py", "body": "x" * 40} + for i in range(20) + ], + } + review = "\n".join(["r" * 24] * 400) + body = append_review_state(review, state, max_chars=10240 - 3) + assert len(body) <= 10240 - 3 + + prepared = provider._prepare_comment_body(body) + + assert len(prepared) <= 10240 + parsed = parse_review_state(prepared) + assert parsed.valid + assert [f["finding_id"] for f in parsed.state["findings"]] == [f"f{i}" for i in range(20)] + assert prepared.rstrip().endswith("-->") + assert "..." in prepared + + def test_class_docstring_survives_the_comment_limit_attribute(self): + assert CodeCommitProvider.__doc__ is not None + assert "CodeCommit" in CodeCommitProvider.__doc__ + + def test_prepare_comment_body_leaves_short_comment_alone(self): + provider = self._make_persistent_provider() + + assert provider._prepare_comment_body("line one\nline two") == "line one\n\nline two" + + def test_publish_comment_sends_capped_body(self): + provider = self._make_persistent_provider() + provider.codecommit_client.publish_comment.return_value = {"comment": {}} + + provider.publish_comment("\n".join(["y" * 100] * 120)) + + sent = provider.codecommit_client.publish_comment.call_args.kwargs["comment"] + assert len(sent) <= 10240 + assert sent.endswith("...") + + def test_edit_comment_sends_capped_body(self): + provider = self._make_persistent_provider() + provider.codecommit_client.update_comment.return_value = {"comment": {}} + + provider.edit_comment({"id": "comment-1"}, "\n".join(["z" * 100] * 120)) + + sent = provider.codecommit_client.update_comment.call_args.args[1] + assert len(sent) <= 10240 + assert sent.endswith("...") + def test_publish_comment_uses_every_pull_request_target(self): provider = object.__new__(CodeCommitProvider) provider.repo_name = "source-repository" @@ -241,6 +383,477 @@ def test_publish_comment_uses_every_pull_request_target(self): ), ] + def test_persistent_review_creates_then_updates_each_pull_request_target(self): + targets = [ + SimpleNamespace( + repository_name="source-repository", + source_commit="source-commit-1", + source_branch="refs/heads/feature", + destination_commit="destination-commit-1", + destination_branch="refs/heads/main", + ), + SimpleNamespace( + repository_name="source-repository", + source_commit="source-commit-2", + source_branch="refs/heads/feature", + destination_commit="destination-commit-2", + destination_branch="refs/heads/release", + ), + ] + provider = self._make_persistent_provider(targets=targets) + provider.codecommit_client.get_comments_for_pull_request.return_value = [] + + def publish_comment(**kwargs): + return { + "comment": self._comment( + kwargs["comment"], + comment_id=f"comment-{kwargs['source_commit']}", + ) + } + + provider.codecommit_client.publish_comment.side_effect = publish_comment + provider.codecommit_client.update_comment.side_effect = self._successful_update_response + header = "## Team Review" + + created = provider.publish_persistent_comment( + f"{header}\n\nfirst review", + initial_header=header, + final_update_message=False, + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=f"{PRReviewHeader.REGULAR.value} 🔍", + ) + first_bodies = { + call.kwargs["source_commit"]: call.kwargs["comment"] + for call in provider.codecommit_client.publish_comment.call_args_list + } + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group( + self._comment(first_bodies["source-commit-1"], comment_id="comment-source-commit-1"), + before_commit_id="destination-commit-1", + after_commit_id="source-commit-1", + ), + self._comment_group( + self._comment(first_bodies["source-commit-2"], comment_id="comment-source-commit-2"), + before_commit_id="destination-commit-2", + after_commit_id="source-commit-2", + ), + ] + provider.codecommit_client.publish_comment.reset_mock() + + updated = provider.publish_persistent_comment( + f"{header}\n\nsecond review", + initial_header=header, + final_update_message=False, + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=f"{PRReviewHeader.REGULAR.value} 🔍", + ) + + assert [comment.id for comment in created] == ["comment-source-commit-1", "comment-source-commit-2"] + assert [comment.id for comment in updated] == ["comment-source-commit-1", "comment-source-commit-2"] + provider.codecommit_client.publish_comment.assert_not_called() + assert provider.codecommit_client.update_comment.call_count == 2 + update_bodies = { + call.args[0]: call.args[1] + for call in provider.codecommit_client.update_comment.call_args_list + } + assert set(update_bodies) == {"comment-source-commit-1", "comment-source-commit-2"} + assert PRReviewIdentity.REGULAR.value in first_bodies["source-commit-1"] + assert PRReviewIdentity.REGULAR.value in update_bodies["comment-source-commit-1"] + assert PRReviewIdentity.REGULAR.value in update_bodies["comment-source-commit-2"] + assert "second review" in update_bodies["comment-source-commit-1"] + assert "Review updated until commit source-commit-1" in update_bodies["comment-source-commit-1"] + assert "Review updated until commit source-commit-2" in update_bodies["comment-source-commit-2"] + + def test_persistent_review_updates_comment_from_older_commit_pair_for_single_target(self): + provider = self._make_persistent_provider() + existing_body = f"## Team Review\n\n{PRReviewIdentity.REGULAR.value}\n\nprevious review" + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group( + self._comment(existing_body, comment_id="older-commit-comment"), + before_commit_id="older-destination-commit", + after_commit_id="older-source-commit", + ) + ] + provider.codecommit_client.update_comment.side_effect = self._successful_update_response + + result = provider.publish_persistent_comment( + "## Team Review\n\nnew review", + initial_header="## Team Review", + final_update_message=False, + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=f"{PRReviewHeader.REGULAR.value} 🔍", + ) + + assert result.id == "older-commit-comment" + provider.codecommit_client.publish_comment.assert_not_called() + comment_id, updated_body = provider.codecommit_client.update_comment.call_args.args + assert comment_id == "older-commit-comment" + assert "Review updated until commit source-commit-1" in updated_body + + def test_persistent_review_falls_back_to_new_comment_when_comment_listing_fails(self): + provider = self._make_persistent_provider() + provider.codecommit_client.get_comments_for_pull_request.side_effect = ValueError("AccessDenied") + provider.codecommit_client.publish_comment.return_value = { + "comment": self._comment("created body", comment_id="fallback-comment") + } + + result = provider.publish_persistent_comment( + "## Team Review\n\nnew review", + initial_header="## Team Review", + final_update_message=False, + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=f"{PRReviewHeader.REGULAR.value} 🔍", + ) + + assert result.id == "fallback-comment" + provider.codecommit_client.update_comment.assert_not_called() + provider.codecommit_client.publish_comment.assert_called_once() + assert PRReviewIdentity.REGULAR.value in provider.codecommit_client.publish_comment.call_args.kwargs["comment"] + + def test_persistent_review_does_not_use_repository_only_match_for_multiple_targets(self): + targets = [ + SimpleNamespace( + repository_name="source-repository", + source_commit="source-commit-1", + source_branch="refs/heads/feature", + destination_commit="destination-commit-1", + destination_branch="refs/heads/main", + ), + SimpleNamespace( + repository_name="source-repository", + source_commit="source-commit-2", + source_branch="refs/heads/feature", + destination_commit="destination-commit-2", + destination_branch="refs/heads/release", + ), + ] + provider = self._make_persistent_provider(targets=targets) + existing_body = f"## Team Review\n\n{PRReviewIdentity.REGULAR.value}\n\nprevious review" + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group( + self._comment(existing_body, comment_id="unmatched-comment"), + before_commit_id="older-destination-commit", + after_commit_id="older-source-commit", + ) + ] + + provider.publish_persistent_comment( + "## Team Review\n\nnew review", + initial_header="## Team Review", + final_update_message=False, + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=f"{PRReviewHeader.REGULAR.value} 🔍", + ) + + provider.codecommit_client.update_comment.assert_not_called() + assert provider.codecommit_client.publish_comment.call_args_list == [ + call( + repo_name="source-repository", + pr_number=321, + destination_commit="destination-commit-1", + source_commit="source-commit-1", + comment=provider.codecommit_client.publish_comment.call_args_list[0].kwargs["comment"], + ), + call( + repo_name="source-repository", + pr_number=321, + destination_commit="destination-commit-2", + source_commit="source-commit-2", + comment=provider.codecommit_client.publish_comment.call_args_list[1].kwargs["comment"], + ), + ] + + def test_persistent_review_does_not_use_ambiguous_destination_match_for_multiple_targets(self): + targets = [ + SimpleNamespace( + repository_name="source-repository", + source_commit="source-commit-a-new", + source_branch="refs/heads/feature-a", + destination_commit="shared-destination-commit", + destination_branch="refs/heads/main", + ), + SimpleNamespace( + repository_name="source-repository", + source_commit="source-commit-b-new", + source_branch="refs/heads/feature-b", + destination_commit="shared-destination-commit", + destination_branch="refs/heads/main", + ), + ] + provider = self._make_persistent_provider(targets=targets) + existing_body = f"## Team Review\n\n{PRReviewIdentity.REGULAR.value}\n\nprevious review" + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group( + self._comment( + existing_body, + comment_id="target-b-old", + created_at=datetime(2024, 1, 2, tzinfo=timezone.utc), + ), + before_commit_id="shared-destination-commit", + after_commit_id="source-commit-b-old", + ), + self._comment_group( + self._comment( + existing_body, + comment_id="target-a-old", + created_at=datetime(2024, 1, 1, tzinfo=timezone.utc), + ), + before_commit_id="shared-destination-commit", + after_commit_id="source-commit-a-old", + ), + ] + + provider.publish_persistent_comment( + "## Team Review\n\nnew review", + initial_header="## Team Review", + final_update_message=False, + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=f"{PRReviewHeader.REGULAR.value} 🔍", + ) + + provider.codecommit_client.update_comment.assert_not_called() + assert [ + call.kwargs["source_commit"] + for call in provider.codecommit_client.publish_comment.call_args_list + ] == ["source-commit-a-new", "source-commit-b-new"] + + @pytest.mark.parametrize("raw_order", [("old", "new"), ("new", "old")]) + def test_get_issue_comments_uses_oldest_first_creation_order_for_answer(self, raw_order): + provider = self._make_persistent_provider() + comments_by_name = { + "old": self._comment( + "older answer", + comment_id="old-comment", + created_at=datetime(2024, 1, 1, tzinfo=timezone.utc), + last_modified_at=datetime(2024, 1, 10, tzinfo=timezone.utc), + ), + "new": self._comment( + "newer answer", + comment_id="new-comment", + created_at=datetime(2024, 1, 2, tzinfo=timezone.utc), + ), + } + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group(*(comments_by_name[name] for name in raw_order)) + ] + + comments = provider.get_issue_comments() + + assert provider.is_supported("get_issue_comments") is True + assert provider.supports_review_comment_identity() is True + assert [comment.id for comment in comments] == ["old-comment", "new-comment"] + assert comments[0].user.login == "arn:aws:iam::123456789012:user/pr-agent" + provider.codecommit_client.get_comments_for_pull_request.assert_called_once_with(321) + + @pytest.mark.parametrize("raw_order", [("q-old", "a-old", "q-new", "a-new"), ("a-new", "q-new", "a-old", "q-old")]) + def test_answer_mode_uses_newest_question_and_answer_after_oldest_first_sort(self, raw_order): + provider = self._make_persistent_provider() + bodies_by_name = { + "q-old": "Questions to better understand the PR:\n\nold question", + "a-old": "/answer old answer", + "q-new": "Questions to better understand the PR:\n\nnew question", + "a-new": "/answer new answer", + } + dates_by_name = { + "q-old": datetime(2024, 1, 1, tzinfo=timezone.utc), + "a-old": datetime(2024, 1, 2, tzinfo=timezone.utc), + "q-new": datetime(2024, 1, 3, tzinfo=timezone.utc), + "a-new": datetime(2024, 1, 4, tzinfo=timezone.utc), + } + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group( + *( + self._comment( + bodies_by_name[name], + comment_id=name, + created_at=dates_by_name[name], + ) + for name in raw_order + ) + ) + ] + reviewer = PRReviewer.__new__(PRReviewer) + reviewer.is_answer = True + reviewer.git_provider = provider + + question, answer = reviewer._get_user_answers() + + assert question == bodies_by_name["q-new"] + assert answer == bodies_by_name["a-new"] + + @pytest.mark.parametrize("raw_order", [("old", "new"), ("new", "old")]) + def test_persistent_review_updates_newest_duplicate_under_both_api_feed_orders(self, raw_order): + provider = self._make_persistent_provider() + provider.codecommit_client.update_comment.side_effect = self._successful_update_response + body = f"## Team Review\n\n{PRReviewIdentity.REGULAR.value}\n\nprevious review" + comments_by_name = { + "old": self._comment( + body, + comment_id="old-comment", + created_at=datetime(2024, 1, 1, tzinfo=timezone.utc), + last_modified_at=datetime(2024, 1, 10, tzinfo=timezone.utc), + ), + "new": self._comment( + body, + comment_id="new-comment", + created_at=datetime(2024, 1, 2, tzinfo=timezone.utc), + ), + } + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group(*(comments_by_name[name] for name in raw_order)) + ] + + provider.publish_persistent_comment( + "## Team Review\n\nnew review", + initial_header="## Team Review", + final_update_message=False, + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=f"{PRReviewHeader.REGULAR.value} 🔍", + ) + + provider.codecommit_client.update_comment.assert_called_once() + assert provider.codecommit_client.update_comment.call_args.args[0] == "new-comment" + provider.codecommit_client.publish_comment.assert_not_called() + + def test_get_issue_comments_filters_deleted_replies_inline_and_malformed_comments(self): + provider = self._make_persistent_provider() + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + "not a group", + self._comment_group(self._comment("valid old", comment_id="valid-old")), + self._comment_group( + self._comment("reply", comment_id="reply", inReplyTo="valid-old"), + self._comment("deleted", comment_id="deleted", deleted=True), + self._comment("missing id", commentId=None), + self._comment("", comment_id="empty-body"), + self._comment("bad timestamp", comment_id="bad-timestamp", creationDate="yesterday"), + ), + self._comment_group( + self._comment("inline", comment_id="inline"), + location={"filePath": "src/app.py", "filePosition": 7}, + ), + self._comment_group( + self._comment( + "valid new", + comment_id="valid-new", + created_at=datetime(2024, 1, 2, tzinfo=timezone.utc), + ) + ), + ] + + comments = provider.get_issue_comments() + + assert [(comment.id, comment.body) for comment in comments] == [ + ("valid-old", "valid old"), + ("valid-new", "valid new"), + ] + + def test_persistent_review_migrates_legacy_heading_to_stable_identity(self): + provider = self._make_persistent_provider() + legacy_body = f"{PRReviewHeader.REGULAR.value} 🔍\n\nprevious review" + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group(self._comment(legacy_body, comment_id="legacy-comment")) + ] + provider.codecommit_client.update_comment.side_effect = self._successful_update_response + + provider.publish_persistent_comment( + "## Custom Review\n\nnew review", + initial_header="## Custom Review", + final_update_message=False, + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=f"{PRReviewHeader.REGULAR.value} 🔍", + ) + + provider.codecommit_client.publish_comment.assert_not_called() + comment_id, updated_body = provider.codecommit_client.update_comment.call_args.args + assert comment_id == "legacy-comment" + assert updated_body.startswith("## Custom Review\n\n") + assert PRReviewIdentity.REGULAR.value in updated_body + + def test_persistent_review_falls_back_to_new_comment_only_for_failed_target_update(self): + targets = [ + SimpleNamespace( + repository_name="source-repository", + source_commit="source-commit-1", + source_branch="refs/heads/feature", + destination_commit="destination-commit-1", + destination_branch="refs/heads/main", + ), + SimpleNamespace( + repository_name="source-repository", + source_commit="source-commit-2", + source_branch="refs/heads/feature", + destination_commit="destination-commit-2", + destination_branch="refs/heads/release", + ), + ] + provider = self._make_persistent_provider(targets=targets) + existing_body = f"## Team Review\n\n{PRReviewIdentity.REGULAR.value}\n\nprevious review" + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group( + self._comment(existing_body, comment_id="comment-target-1"), + before_commit_id="destination-commit-1", + after_commit_id="source-commit-1", + ), + self._comment_group( + self._comment(existing_body, comment_id="comment-target-2"), + before_commit_id="destination-commit-2", + after_commit_id="source-commit-2", + ), + ] + + def update_comment(comment_id, body): + if comment_id == "comment-target-2": + return {"comment": {"commentId": comment_id, "content": "unexpected body", "deleted": False}} + return self._successful_update_response(comment_id, body) + + provider.codecommit_client.update_comment.side_effect = update_comment + provider.codecommit_client.publish_comment.return_value = { + "comment": self._comment("fallback body", comment_id="fallback-comment") + } + + result = provider.publish_persistent_comment( + "## Team Review\n\nnew review", + initial_header="## Team Review", + final_update_message=False, + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=f"{PRReviewHeader.REGULAR.value} 🔍", + ) + + assert [comment.id for comment in result] == ["comment-target-1", "fallback-comment"] + assert provider.codecommit_client.update_comment.call_count == 2 + provider.codecommit_client.publish_comment.assert_called_once() + assert provider.codecommit_client.publish_comment.call_args.kwargs["repo_name"] == "source-repository" + assert provider.codecommit_client.publish_comment.call_args.kwargs["destination_commit"] == "destination-commit-2" + assert provider.codecommit_client.publish_comment.call_args.kwargs["source_commit"] == "source-commit-2" + + def test_unmarked_persistent_comment_remains_create_only(self): + provider = self._make_persistent_provider() + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + self._comment_group(self._comment("## Title\n\nold description", comment_id="describe-comment")) + ] + provider.codecommit_client.publish_comment.return_value = { + "comment": self._comment("## Title\n\nnew description", comment_id="new-describe-comment") + } + + result = provider.publish_persistent_comment( + "## Title\n\nnew description", + initial_header="## Title", + update_header=True, + name="describe", + final_update_message=False, + ) + + assert result.id == "new-describe-comment" + provider.codecommit_client.get_comments_for_pull_request.assert_not_called() + provider.codecommit_client.update_comment.assert_not_called() + provider.codecommit_client.publish_comment.assert_called_once() + + def test_get_issue_comments_support_does_not_enable_gfm_improve_history(self): + provider = self._make_persistent_provider() + + assert provider.is_supported("get_issue_comments") is True + assert provider.is_supported("gfm_markdown") is False + def test_publish_code_suggestion_uses_target_that_contains_file(self): provider = object.__new__(CodeCommitProvider) provider.repo_name = "source-repository" diff --git a/tests/unittest/test_command_descriptions.py b/tests/unittest/test_command_descriptions.py index 4e67007fa3..32cd3c8519 100644 --- a/tests/unittest/test_command_descriptions.py +++ b/tests/unittest/test_command_descriptions.py @@ -9,6 +9,8 @@ from pr_agent.tools.pr_help_message import PRHelpMessage REPO_ROOT = Path(__file__).resolve().parents[2] +CURRENT_DOCS_URL = "https://docs.pr-agent.ai" +RETIRED_DOCS_HOSTS = ("qodo-merge-docs.qodo.ai", "pr-agent-docs.codium.ai") def test_cli_and_bot_help_use_canonical_command_descriptions(): @@ -32,6 +34,24 @@ def test_tool_usage_guides_use_canonical_command_descriptions(usage_guide, comma assert COMMAND_DESCRIPTIONS[command] in usage_guide() +def test_runtime_help_uses_current_documentation_links(): + help_text = "\n".join( + [ + HelpMessage.get_general_commands_text(), + HelpMessage.get_general_bot_help_text(), + HelpMessage.get_review_usage_guide(), + HelpMessage.get_describe_usage_guide(), + HelpMessage.get_ask_usage_guide(), + HelpMessage.get_improve_usage_guide(), + HelpMessage.get_help_docs_usage_guide(), + ] + ) + + assert f"{CURRENT_DOCS_URL}/tools/review/#configuration-options" in help_text + assert f"{CURRENT_DOCS_URL}/tools/ask/" in help_text + assert not any(host in help_text for host in RETIRED_DOCS_HOSTS) + + @pytest.mark.parametrize("command", COMMAND_DESCRIPTIONS) def test_tool_docs_use_canonical_command_descriptions(command): tool_docs = REPO_ROOT / "docs" / "docs" / "tools" / f"{command}.md" @@ -48,6 +68,14 @@ def __init__(self): def is_supported(_feature): return True + @staticmethod + def supports_checkbox_commands(): + return False + + @staticmethod + def supports_markdown_tables(): + return False + def publish_comment(self, comment): self.comment = comment diff --git a/tests/unittest/test_config_dead_keys.py b/tests/unittest/test_config_dead_keys.py new file mode 100644 index 0000000000..ba37c7f67c --- /dev/null +++ b/tests/unittest/test_config_dead_keys.py @@ -0,0 +1,219 @@ +"""Regression guard for dead configuration keys. + +Every key in pr_agent/settings/configuration.toml must be read somewhere in +pr_agent/ source code. Static detection is deliberately conservative: keys are +matched through the common read patterns below, and keys that are genuinely +used through variable indirection (e.g. a settings section assigned to a local +variable and then indexed) are recorded in the allowlist with a note pointing +at the reading site. + +A whole-section read (settings.get("SECTION", ...)) does not count as reading +its keys: each key must be matched through an attribute chain, a quoted +dotted path, or a section-level .get("KEY") / getattr() call, or it must be +allowlisted. + +Adding a key without a reader - or leaving a config-only key in the TOML - +fails this test, so the "dead key ledger" cannot silently regrow. +""" + +import re +import tomllib +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +CONFIG_TOML = ROOT / "pr_agent/settings/configuration.toml" +PR_AGENT_SOURCE = ROOT / "pr_agent" + +# Keys read through a settings section stored in a local variable, so the +# static dotted-path scan cannot see them. Each entry documents the reader. +_ALLOWLIST = { + # [artifacts] read as artifacts_settings = get_settings().get("ARTIFACTS", {}) + ( + "artifacts", + "max_artifact_size", + ): "artifacts_settings.get('max_artifact_size', 50000) in pr_agent/algo/artifacts.py", + ( + "artifacts", + "artifact_label", + ): "artifacts_settings.get('artifact_label', '') in pr_agent/algo/artifacts.py", + # [config] helpers that re-index the section from a variable or expression + ( + "config", + "reaction_on_failure", + ): "get_reaction_setting('reaction_on_success' if succeeded else 'reaction_on_failure') " + "reads get_settings().config.get(name) in pr_agent/git_providers/git_provider.py", + ( + "config", + "output_relevant_configurations", + ): "get_settings().get('config', {}).get('output_relevant_configurations', False) in " + "pr_agent/tools/pr_description.py", + ( + "config", + "enable_claude_adaptive_thinking", + ): "settings.config.get(key, default) in the thinking-controls loop in " + "pr_agent/algo/ai_handlers/litellm_ai_handler.py", + ( + "config", + "enable_claude_extended_thinking", + ): "settings.config.get(key, default) in the thinking-controls loop in " + "pr_agent/algo/ai_handlers/litellm_ai_handler.py", + ( + "config", + "extended_thinking_budget_tokens", + ): "settings.config.get(key, default) in the thinking-controls loop in " + "pr_agent/algo/ai_handlers/litellm_ai_handler.py", + ( + "config", + "extended_thinking_max_output_tokens", + ): "settings.config.get(key, default) in the thinking-controls loop in " + "pr_agent/algo/ai_handlers/litellm_ai_handler.py", + # [gerrit] read as gerrit_settings = get_settings().get("gerrit", {}) + ( + "gerrit", + "webhook_username", + ): "gerrit_settings.get('webhook_username', None) in pr_agent/servers/gerrit_server.py", + ( + "gerrit", + "webhook_password", + ): "gerrit_settings.get('webhook_password', None) in pr_agent/servers/gerrit_server.py", + ( + "github", + "api_retries", + ): "github_config = get_settings().github then github_config.get('api_retries') in " + "pr_agent/git_providers/github_provider.py", + ( + "github", + "seconds_between_requests", + ): "github_config.get('seconds_between_requests') in pr_agent/git_providers/github_provider.py", + ( + "github", + "seconds_between_writes", + ): "github_config.get('seconds_between_writes') in pr_agent/git_providers/github_provider.py", + # [openrouter] read as openrouter_settings = settings.get("openrouter", {}) or {} + ( + "openrouter", + "provider_only", + ): "openrouter_settings.get('provider_only', []) in pr_agent/algo/ai_handlers/litellm_ai_handler.py", + ( + "openrouter", + "provider_order", + ): "openrouter_settings.get('provider_order', []) in pr_agent/algo/ai_handlers/litellm_ai_handler.py", + ( + "openrouter", + "allow_fallbacks", + ): "openrouter_settings.get('allow_fallbacks', True) in pr_agent/algo/ai_handlers/litellm_ai_handler.py", + ( + "openrouter", + "reasoning_effort", + ): "openrouter_settings.get('reasoning_effort', '') in pr_agent/algo/ai_handlers/litellm_ai_handler.py", + ( + "openrouter", + "reasoning_max_tokens", + ): "openrouter_settings.get('reasoning_max_tokens', 0) in pr_agent/algo/ai_handlers/litellm_ai_handler.py", + ( + "openrouter", + "max_tokens", + ): "openrouter_settings.get('max_tokens', 0) in pr_agent/algo/ai_handlers/litellm_ai_handler.py", + ( + "pr_description", + "pr_diagram_direction", + ): "description_settings = get_settings().pr_description then " + "description_settings.pr_diagram_direction in pr_agent/tools/pr_description.py", + ( + "pr_description", + "pr_diagram_direction_threshold", + ): "description_settings.pr_diagram_direction_threshold in pr_agent/tools/pr_description.py", + # [push_outputs] read as cfg = get_settings().get('push_outputs', {}) or {} + ( + "push_outputs", + "enable", + ): "cfg.get('enable', False) in pr_agent/algo/utils.py", + ( + "push_outputs", + "channels", + ): "cfg.get('channels', []) in pr_agent/algo/utils.py", + ( + "push_outputs", + "file_path", + ): "cfg.get('file_path', 'pr-agent-outputs/reviews.jsonl') in pr_agent/algo/utils.py", + ( + "push_outputs", + "webhook_url", + ): "_push_outputs_sink_url(cfg, 'webhook_url') in pr_agent/algo/utils.py", + ( + "push_outputs", + "slack_webhook_url", + ): "_push_outputs_sink_url(cfg, 'slack_webhook_url') in pr_agent/algo/utils.py", +} + + +def _read_sources() -> str: + return "\n".join( + p.read_text(encoding="utf-8") for p in sorted(PR_AGENT_SOURCE.rglob("*.py")) + ).lower() + + +def _key_is_read(section: str, key: str, sources: str) -> bool: + section_name = re.escape(section) + key_name = re.escape(key) + settings_prefix = ( + r"(?:get_settings\s*\(\s*\)|global_settings|self\s*\.\s*settings|settings)" + ) + # 1. attribute chain: settings.
      . + if re.search( + settings_prefix + r"\s*\.\s*" + section_name + r"\s*\.\s*" + key_name + r"\b", + sources, + ): + return True + # 2. quoted dotted paths: "SECTION.KEY" or "section.key" + for dotted in (f"{section.upper()}.{key.upper()}", f"{section}.{key}"): + if re.search(re.escape(f"\"{dotted}\"") + r"|" + re.escape(f"'{dotted}'"), sources): + return True + # 3. settings.
      .get("KEY") + if re.search( + settings_prefix + r"\s*\.\s*" + section_name + r"\s*\.\s*get\s*\(\s*['\"]" + + key_name + r"['\"]", + sources, + ): + return True + # 4. getattr(settings.
      , "KEY", ...) + if re.search( + r"getattr\s*\(\s*" + settings_prefix + r"\s*\.\s*" + section_name + + r"\s*,\s*['\"]" + key_name + r"['\"]", + sources, + ): + return True + # 5. helper indirection for the [config] section: _read_bool_setting("key") + # and get_reaction_setting("key") both read get_settings().config.get(key). + if section == "config": + for helper in ("_read_bool_setting", "get_reaction_setting"): + if re.search( + r"\b" + re.escape(helper) + r"\s*\(\s*['\"]" + key_name + r"['\"]", + sources, + ): + return True + return False + + +def test_every_config_key_has_a_reader_or_is_allowlisted(): + sources = _read_sources() + with open(CONFIG_TOML, "rb") as f: + settings = tomllib.load(f) + + unread = { + (section, key) + for section, entries in settings.items() + if not section.startswith("_") + for key in entries + if not _key_is_read(section, key, sources) + } + allowlisted = set(_ALLOWLIST) + + assert unread <= allowlisted, ( + "configuration.toml keys have no detected reader; remove them or allowlist them: " + f"{sorted(unread - allowlisted)}" + ) + assert ( + allowlisted - unread == set() + ), "allowlist has stale entries (those keys are now detected as read or no longer present): " \ + f"{sorted(allowlisted - unread)}" diff --git a/tests/unittest/test_config_reference.py b/tests/unittest/test_config_reference.py new file mode 100644 index 0000000000..3d625afff3 --- /dev/null +++ b/tests/unittest/test_config_reference.py @@ -0,0 +1,43 @@ +import re +from collections import Counter +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +CONFIG_TOML = ROOT / "pr_agent/settings/configuration.toml" +OUTPUT_PAGE = ROOT / "docs/docs/usage-guide/configuration_reference.md" + + +def _toml_keys(text: str) -> Counter: + return Counter(re.findall(r"^([a-zA-Z0-9_]+)\s*=", text, re.M)) + + +def _documented_keys(text: str) -> Counter: + return Counter( + row.group(1) + for line in text.splitlines() + if line.startswith("| ") + for row in [re.match(r"\| `([^`]+)` \|", line)] + if row + ) + + +def test_config_reference_covers_every_active_key(): + keys = _toml_keys(CONFIG_TOML.read_text(encoding="utf-8")) + assert sum(keys.values()) == 262 + + assert { + "model", + "enable_auto_approval", + "reaction_on_start", + "reaction_on_failure", + } <= set(keys) + assert {"force_streaming_custom_llm_provider", "cache_control_injection_points"} <= set(keys) + + +def test_config_reference_page_lists_every_active_key(): + toml_keys = _toml_keys(CONFIG_TOML.read_text(encoding="utf-8")) + documented = _documented_keys(OUTPUT_PAGE.read_text(encoding="utf-8")) + assert documented == toml_keys, ( + "docs/docs/usage-guide/configuration_reference.md is out of date; " + "run scripts/generate_config_reference.py" + ) diff --git a/tests/unittest/test_docs_nav_parity.py b/tests/unittest/test_docs_nav_parity.py new file mode 100644 index 0000000000..33b3ee8578 --- /dev/null +++ b/tests/unittest/test_docs_nav_parity.py @@ -0,0 +1,72 @@ +"""Guards for issue #3198: keep the two documentation navigations in step. + +The docs are rendered twice. `docs/mkdocs.yml` drives the mkdocs site, and +`docs/docs/summary.md` is the GitBook table of contents, wired up in +`docs/docs/.gitbook.yaml` as `summary: ./summary.md`. `extending_pr_agent.md` +tells contributors to register every new page in both, so the two files are +expected to list exactly the same set of pages — and nothing enforces that +today. These tests do. +""" + +import re +from pathlib import Path + +DOCS = Path(__file__).resolve().parents[2] / "docs" +PAGES = DOCS / "docs" +MKDOCS = DOCS / "mkdocs.yml" +SUMMARY = PAGES / "summary.md" + +# Pages that are deliberately outside both navigations. +NOT_IN_NAV = { + # summary.md is a navigation file itself, not a page of the site. + "summary.md", +} + + +def _mkdocs_pages() -> set[str]: + """Page paths listed in the mkdocs nav, ignoring commented-out entries.""" + lines = [ + line + for line in MKDOCS.read_text(encoding="utf-8").splitlines() + if not line.lstrip().startswith("#") + ] + return set(re.findall(r"'([^']+\.md)'", "\n".join(lines))) + + +def _summary_pages() -> set[str]: + """Page paths linked from the GitBook summary.""" + return set(re.findall(r"\(([^()]+\.md)\)", SUMMARY.read_text(encoding="utf-8"))) + + +def test_every_page_is_reachable_from_a_nav(): + """No page under docs/docs is orphaned from both navigations.""" + on_disk = {p.relative_to(PAGES).as_posix() for p in PAGES.rglob("*.md")} + reachable = _mkdocs_pages() | _summary_pages() | NOT_IN_NAV + assert not on_disk - reachable, ( + "pages exist but no navigation points at them: " + f"{sorted(on_disk - reachable)}" + ) + + +def test_navs_list_the_same_pages(): + """mkdocs.yml and summary.md must not drift apart.""" + mkdocs, summary = _mkdocs_pages(), _summary_pages() + assert not mkdocs - summary, ( + f"in mkdocs.yml but missing from summary.md: {sorted(mkdocs - summary)}" + ) + assert not summary - mkdocs, ( + f"in summary.md but missing from mkdocs.yml: {sorted(summary - mkdocs)}" + ) + + +def test_navs_do_not_point_at_missing_files(): + """Every navigation entry resolves to a file that exists.""" + for label, pages in (("mkdocs.yml", _mkdocs_pages()), ("summary.md", _summary_pages())): + missing = sorted(page for page in pages if not (PAGES / page).exists()) + assert not missing, f"{label} references missing pages: {missing}" + + +def test_declared_exclusions_still_exist(): + """Keeps NOT_IN_NAV honest: a stale entry there would hide a real orphan.""" + missing = sorted(name for name in NOT_IN_NAV if not (PAGES / name).exists()) + assert not missing, f"NOT_IN_NAV lists files that no longer exist: {missing}" diff --git a/tests/unittest/test_effort_bar_clamp.py b/tests/unittest/test_effort_bar_clamp.py index b11484b3c4..a411e66ac3 100644 --- a/tests/unittest/test_effort_bar_clamp.py +++ b/tests/unittest/test_effort_bar_clamp.py @@ -40,7 +40,7 @@ def test_the_bar_is_always_five_segments(): @pytest.mark.parametrize("score, expected", [("8", 5), ("0", 1), ("3", 3)]) -def test_the_review_label_matches_the_rendered_bar(score, expected): +def test_the_review_label_matches_the_rendered_bar(score, expected, monkeypatch): """Clamp the label the same way as the bar, so the two never disagree.""" from pr_agent.config_loader import get_settings from pr_agent.tools.pr_reviewer import PRReviewer @@ -54,11 +54,11 @@ def test_the_review_label_matches_the_rendered_bar(score, expected): })() settings = get_settings(use_context=False) - settings.set("pr_reviewer.enable_review_labels_effort", True) - settings.set("pr_reviewer.enable_review_labels_security", False) - settings.set("pr_reviewer.require_estimate_effort_to_review", True) - settings.set("pr_reviewer.require_security_review", False) - settings.set("config.publish_output", True) + monkeypatch.setattr(settings.pr_reviewer, "enable_review_labels_effort", True) + monkeypatch.setattr(settings.pr_reviewer, "enable_review_labels_security", False) + monkeypatch.setattr(settings.pr_reviewer, "require_estimate_effort_to_review", True) + monkeypatch.setattr(settings.pr_reviewer, "require_security_review", False) + monkeypatch.setattr(settings.config, "publish_output", True) reviewer.set_review_labels({"review": {"estimated_effort_to_review_[1-5]": score}}) assert published and published[0] == [f"Review effort {expected}/5"] diff --git a/tests/unittest/test_find_line_number_of_relevant_line_in_file.py b/tests/unittest/test_find_line_number_of_relevant_line_in_file.py index 06db14f08d..18e1eb72e4 100644 --- a/tests/unittest/test_find_line_number_of_relevant_line_in_file.py +++ b/tests/unittest/test_find_line_number_of_relevant_line_in_file.py @@ -90,3 +90,54 @@ def test_hunk_header_with_omitted_count_on_the_new_side_only(self): assert position == 1 assert absolute_position == 1 + + def test_combined_diff_header_does_not_crash(self): + """Skip a combined/merge hunk header (@@@) instead of crashing on its None regex match.""" + diff_files = [ + FilePatchInfo(base_file="file1", head_file="file1", + patch="@@@ -1,2 -1,2 +1,4 @@@\n- a\n +b\n c\n +d", filename="file1") + ] + + position, absolute_position = find_line_number_of_relevant_line_in_file( + diff_files, "file1", "c") + + assert position == -1 + assert absolute_position == -1 + + def test_combined_diff_header_absolute_position_does_not_crash(self): + diff_files = [ + FilePatchInfo(base_file="file1", head_file="file1", + patch="@@@ -1,2 -1,2 +1,4 @@@\n- a\n +b\n c\n +d", filename="file1") + ] + + position, absolute_position = find_line_number_of_relevant_line_in_file( + diff_files, "file1", "not_found", absolute_position=3) + + assert position == -1 + + def test_combined_diff_header_plus_line_does_not_crash(self): + diff_files = [ + FilePatchInfo(base_file="file1", head_file="file1", + patch="@@@ -1,2 -1,2 +1,4 @@@\n- a\n +b\n c\n +d", filename="file1") + ] + + position, absolute_position = find_line_number_of_relevant_line_in_file( + diff_files, "file1", "+b") + + assert position == -1 + assert absolute_position == -1 + + def test_a_valid_hunk_after_a_combined_header_is_still_found(self): + """Keep anchoring from a valid hunk rendered after a skipped combined header.""" + combined = "@@@ -1,2 -1,2 +1,4 @@@\n- a\n +b\n c\n +d" + valid_hunk = "@@ -3,1 +5,1 @@\n ctx\n-old\n+new" + diff_files = [ + FilePatchInfo(base_file="file1", head_file="file1", + patch=combined + "\n" + valid_hunk, filename="file1") + ] + + position, absolute_position = find_line_number_of_relevant_line_in_file( + diff_files, "file1", "new") + + assert position == 8 + assert absolute_position == 6 diff --git a/tests/unittest/test_get_max_tokens.py b/tests/unittest/test_get_max_tokens.py index 7525fc3bc7..203942bbfb 100644 --- a/tests/unittest/test_get_max_tokens.py +++ b/tests/unittest/test_get_max_tokens.py @@ -1058,3 +1058,70 @@ def test_claude_registries_baseline_parity_and_no_duplicates(self): claude_thinking = {m for m in CLAUDE_EXTENDED_THINKING_MODELS if "claude" in m} claude_no_temp = {m for m in NO_SUPPORT_TEMPERATURE_MODELS if "claude" in m} assert claude_thinking.isdisjoint(claude_no_temp) + + @pytest.mark.parametrize( + "model, expected", + [ + ("gpt-4o", 128000), + ("gpt-4.1", 1047576), + ], + ) + def test_ignore_max_model_tokens_returns_unreduced_litellm_value(self, monkeypatch, model, expected): + """Sites that bypass the max_model_tokens clamp (pr_help_message, pr_help_docs, + pr_code_suggestions) must still get the raw model context even after the exact + LiteLLM duplicates left the registry: the value now resolves from LiteLLM, while + the default path keeps clamping to config.max_model_tokens.""" + fake_settings = type('', (), { + 'config': type('', (), { + 'custom_model_max_tokens': 0, + 'max_model_tokens': 32000 + })() + })() + + monkeypatch.setattr(utils, "get_settings", lambda: fake_settings) + monkeypatch.setattr(utils, "MAX_TOKENS", {}) # simulate deletion of the entry + + assert get_max_tokens(model) == 32000 + assert get_max_tokens(model, ignore_max_model_tokens=True) == expected + + +class TestNoLiteLLMDuplicates: + + # Models pinned in MAX_TOKENS because LiteLLM's bundled backup cost map (used + # when the import-time fetch fails, and under LITELLM_LOCAL_MODEL_COST_MAP=true) + # does not carry them, so the get_max_tokens() fallback cannot resolve them. + # They are exempt from the no-duplicates guard by design. + LITELLM_BUNDLED_MAP_UNKNOWN = { + "gemini/gemini-3.8-flash", + "vertex_ai/gemini-3.8-flash", + "openrouter/x-ai/grok-4.5", + "openrouter/x-ai/grok-4.6", + "xai/grok-build-latest", + } + + def test_static_max_tokens_has_no_exact_litellm_duplicates(self): + """Hardcoded MAX_TOKENS entries must not just mirror LiteLLM. + + get_max_tokens() already falls back to litellm.get_model_info(), so a + static entry that reports the identical value is dead duplication. + Generator-expanded Claude families are excluded: they also drive the + no-temperature / extended-thinking registries, and their 1M-context + handling is a separate, deliberate judgement (issue #3196). Entries in + LITELLM_BUNDLED_MAP_UNKNOWN are pinned because the bundled cost map does + not carry them, so the fallback could not resolve them. + """ + generated = set(_generate_claude_registries()[0]) + static = { + k: v + for k, v in MAX_TOKENS.items() + if k not in generated and k not in self.LITELLM_BUNDLED_MAP_UNKNOWN + } + dups = [] + for model, ours in static.items(): + try: + theirs = int(litellm.get_model_info(model).get("max_input_tokens")) + except Exception: + continue + if theirs == ours: + dups.append((model, ours)) + assert not dups, f"MAX_TOKENS entries that exactly duplicate LiteLLM: {sorted(dups)}" diff --git a/tests/unittest/test_git_provider_method_contract.py b/tests/unittest/test_git_provider_method_contract.py index ab7ae4d65d..ee15a5ecb9 100644 --- a/tests/unittest/test_git_provider_method_contract.py +++ b/tests/unittest/test_git_provider_method_contract.py @@ -8,6 +8,7 @@ import inspect from collections.abc import Callable from dataclasses import dataclass +from datetime import datetime, timezone from enum import Enum from types import SimpleNamespace from typing import get_type_hints @@ -53,6 +54,28 @@ class MethodContract: check_execution: bool = True +@dataclass(frozen=True) +class DeliberateMismatch: + reason: str + + +@dataclass(frozen=True) +class PredicateContract: + name: str + evidence: tuple[str, ...] + deliberate_mismatches: dict[str, DeliberateMismatch] + + +@dataclass(frozen=True) +class SuggestionOutcomeContract: + provider_name: str + build_provider: Callable[..., GitProvider] + make_fail: Callable[..., None] + make_succeed: Callable[..., None] + payload: list[dict] + deliberate_mismatch: DeliberateMismatch | None = None + + def _github(monkeypatch) -> GithubProvider: provider = GithubProvider.__new__(GithubProvider) provider.base_url = "https://api.github.example" @@ -125,6 +148,7 @@ def _bitbucket_server(monkeypatch) -> BitbucketServerProvider: }] return provider + def _bitbucket(monkeypatch) -> BitbucketProvider: provider = BitbucketProvider.__new__(BitbucketProvider) provider.pr = MagicMock() @@ -136,6 +160,38 @@ def _bitbucket(monkeypatch) -> BitbucketProvider: return provider +def _codecommit(monkeypatch) -> CodeCommitProvider: + provider = CodeCommitProvider.__new__(CodeCommitProvider) + provider.repo_name = "repo" + provider.pr_num = 7 + provider.pr_url = "https://us-east-1.console.aws.amazon.com/codesuite/codecommit/repositories/repo/pull-requests/7" + provider.pr = SimpleNamespace( + source_commit="source", + destination_commit="destination", + targets=[ + SimpleNamespace( + repository_name="repo", + source_commit="source", + destination_commit="destination", + ) + ], + ) + provider.codecommit_client = MagicMock() + provider.codecommit_client.get_comments_for_pull_request.return_value = [ + { + "repositoryName": "repo", + "beforeCommitId": "destination", + "afterCommitId": "source", + "comments": [{ + "commentId": "comment-1", + "content": COMMENT_BODY, + "creationDate": datetime(2024, 1, 1, tzinfo=timezone.utc), + }], + } + ] + return provider + + def _bare(provider_type): """A provider with no backend wired at all: every call on it must succeed without one.""" return lambda monkeypatch: provider_type.__new__(provider_type) @@ -149,7 +205,7 @@ def _bare(provider_type): "azure-devops": (AzureDevopsProvider, _azure_devops), "bitbucket": (BitbucketProvider, _bitbucket), "bitbucket-server": (BitbucketServerProvider, _bitbucket_server), - "codecommit": (CodeCommitProvider, _bare(CodeCommitProvider)), + "codecommit": (CodeCommitProvider, _codecommit), "local": (LocalGitProvider, _bare(LocalGitProvider)), "plain-diff": (PlainDiffGitProvider, _bare(PlainDiffGitProvider)), "mosaico-diff": (DiffInputProvider, _bare(DiffInputProvider)), @@ -183,6 +239,30 @@ def _is_success(value): REACTION_TIERS = _tiers(supported=("github", "gitlab", "gitea"), not_implemented=("gerrit",)) +PREDICATE_CONTRACTS = ( + PredicateContract( + name="supports_review_comment_identity", + evidence=("edit_comment",), + deliberate_mismatches={ + "gitea": DeliberateMismatch( + "Gitea forwards identity arguments but cannot safely activate identity tracking " + "until it normalizes dictionary-shaped comment payloads." + ), + }, + ), + PredicateContract( + name="supports_thread_resolution", + evidence=("resolve_comment_thread",), + deliberate_mismatches={ + "gitlab": DeliberateMismatch( + "GitLab resolves note IDs while /ask_line addresses discussion IDs, so thread resolution " + "must remain disabled." + ), + }, + ), +) + + METHOD_CONTRACTS = ( MethodContract( name="get_commit_messages", @@ -197,8 +277,17 @@ def _is_success(value): noop_value=[], check_supported=_is_comment_sequence, tiers=_tiers( - supported=("github", "gitlab", "gitea", "gerrit", "azure-devops", "bitbucket-server", "bitbucket"), - not_implemented=("codecommit", "local"), + supported=( + "github", + "gitlab", + "gitea", + "gerrit", + "azure-devops", + "bitbucket-server", + "bitbucket", + "codecommit", + ), + not_implemented=("local",), ), # Implementations narrow the base `Iterable` (a paginated list, a list of SDK objects), # so the return annotation is checked by behaviour rather than by equality. @@ -229,6 +318,187 @@ def _is_success(value): check_return_annotation=False, check_execution=False, ), + MethodContract( + name="get_repo_context_ref", + args=(), + noop_value=None, + check_supported=lambda _: None, + tiers=_tiers( + supported=("github", "gitlab", "gitea", "azure-devops", "bitbucket", "bitbucket-server"), + ), + # Signature + return-annotation contract: the providers that fetch repo-context files + # must expose the same hook so the cache can key on the revision being read. + check_execution=False, + ), +) + + +SUGGESTION_PAYLOAD = [{ + "body": "description\n```suggestion\nnew\n```", + "relevant_file": "app.py", + "relevant_lines_start": 1, + "relevant_lines_end": 1, +}] + + +def _build_github_suggestion_provider(monkeypatch, tmp_path) -> GithubProvider: + provider = _github(monkeypatch) + provider.validate_comments_inside_hunks = lambda suggestions: suggestions + return provider + + +def _build_gitlab_suggestion_provider(monkeypatch, tmp_path) -> GitLabProvider: + provider = _gitlab(monkeypatch) + provider.resolve_outdated_inline_threads = MagicMock() + provider.get_diff_files = MagicMock(return_value=[SimpleNamespace(filename="app.py", head_file="orig\n")]) + return provider + + +def _build_gerrit_suggestion_provider(monkeypatch, tmp_path) -> GerritProvider: + provider = _gerrit(monkeypatch) + provider.repo_path = str(tmp_path) + (tmp_path / "app.py").write_text("orig\n") + monkeypatch.setattr("pr_agent.git_providers.gerrit_provider.upload_patch", lambda *_: "https://patch.example/1") + monkeypatch.setattr("pr_agent.git_providers.gerrit_provider.diff", lambda *_, **__: "patch") + monkeypatch.setattr("pr_agent.git_providers.gerrit_provider.reset_local_changes", lambda *_: None) + return provider + + +def _build_azure_devops_suggestion_provider(monkeypatch, tmp_path) -> AzureDevopsProvider: + provider = _azure_devops(monkeypatch) + provider.workspace_slug = "project" + provider.repo_slug = "repo" + provider.pr_num = 7 + provider.azure_devops_client = MagicMock() + provider._resolve_diff_file_path = MagicMock(return_value="/app.py") + provider._get_suggestion_end_offset = MagicMock(return_value=1) + return provider + + +def _build_codecommit_suggestion_provider(monkeypatch, tmp_path) -> CodeCommitProvider: + provider = _codecommit(monkeypatch) + provider.pr_num = 123 + provider.codecommit_client = MagicMock() + return provider + + +def _succeed_codecommit_suggestions(provider: CodeCommitProvider, monkeypatch, tmp_path): + provider._get_target_contexts_for_file = MagicMock(return_value=[{ + "repository_name": "repo", + "destination_commit": "dest", + "source_commit": "src", + }]) + + +def _build_local_suggestion_provider(monkeypatch, tmp_path) -> LocalGitProvider: + provider = _bare(LocalGitProvider)(monkeypatch) + provider.improve_path = str(tmp_path / "improve.md") + return provider + + +def _build_plain_diff_suggestion_provider(monkeypatch, tmp_path) -> PlainDiffGitProvider: + provider = _bare(PlainDiffGitProvider)(monkeypatch) + provider.output_path = None + return provider + + +SUGGESTION_OUTCOME_CONTRACTS = ( + SuggestionOutcomeContract( + provider_name="github", + build_provider=_build_github_suggestion_provider, + make_fail=lambda p, mp, tmp: setattr(p, "publish_inline_comments", MagicMock(return_value=False)), + make_succeed=lambda p, mp, tmp: setattr(p, "publish_inline_comments", MagicMock(return_value=True)), + payload=SUGGESTION_PAYLOAD, + ), + SuggestionOutcomeContract( + provider_name="gitlab", + build_provider=_build_gitlab_suggestion_provider, + make_fail=lambda p, mp, tmp: setattr( + p, "send_inline_comment", MagicMock(side_effect=RuntimeError("network down")) + ), + make_succeed=lambda p, mp, tmp: setattr(p, "send_inline_comment", MagicMock(return_value=True)), + payload=SUGGESTION_PAYLOAD, + deliberate_mismatch=DeliberateMismatch( + "GitLab unconditionally returns True; issue #3129 owns reporting total failures." + ), + ), + SuggestionOutcomeContract( + provider_name="gitea", + build_provider=lambda mp, tmp: _gitea(mp), + make_fail=lambda p, mp, tmp: setattr(p, "publish_inline_comments", MagicMock(return_value=False)), + make_succeed=lambda p, mp, tmp: setattr(p, "publish_inline_comments", MagicMock(return_value=True)), + payload=SUGGESTION_PAYLOAD, + ), + SuggestionOutcomeContract( + provider_name="gerrit", + build_provider=_build_gerrit_suggestion_provider, + make_fail=lambda p, mp, tmp: mp.setattr( + "pr_agent.git_providers.gerrit_provider.add_comment", + MagicMock(side_effect=RuntimeError("network down")), + ), + make_succeed=lambda p, mp, tmp: mp.setattr( + "pr_agent.git_providers.gerrit_provider.add_comment", lambda *_: None + ), + payload=SUGGESTION_PAYLOAD, + ), + SuggestionOutcomeContract( + provider_name="azure-devops", + build_provider=_build_azure_devops_suggestion_provider, + make_fail=lambda p, mp, tmp: setattr( + p.azure_devops_client, "create_thread", MagicMock(side_effect=RuntimeError("network down")) + ), + make_succeed=lambda p, mp, tmp: setattr(p.azure_devops_client, "create_thread", MagicMock()), + payload=SUGGESTION_PAYLOAD, + ), + SuggestionOutcomeContract( + provider_name="bitbucket", + build_provider=lambda mp, tmp: _bitbucket(mp), + make_fail=lambda p, mp, tmp: setattr(p, "publish_inline_comments", MagicMock(return_value=False)), + make_succeed=lambda p, mp, tmp: setattr(p, "publish_inline_comments", MagicMock(return_value=True)), + payload=SUGGESTION_PAYLOAD, + ), + SuggestionOutcomeContract( + provider_name="bitbucket-server", + build_provider=lambda mp, tmp: _bitbucket_server(mp), + make_fail=lambda p, mp, tmp: setattr(p, "publish_inline_comments", MagicMock(return_value=False)), + make_succeed=lambda p, mp, tmp: setattr(p, "publish_inline_comments", MagicMock(return_value=True)), + payload=SUGGESTION_PAYLOAD, + ), + SuggestionOutcomeContract( + provider_name="codecommit", + build_provider=_build_codecommit_suggestion_provider, + make_fail=lambda p, mp, tmp: setattr(p, "_get_target_contexts_for_file", MagicMock(return_value=[])), + make_succeed=_succeed_codecommit_suggestions, + payload=SUGGESTION_PAYLOAD, + ), + SuggestionOutcomeContract( + provider_name="local", + build_provider=_build_local_suggestion_provider, + make_fail=lambda p, mp, tmp: None, + make_succeed=lambda p, mp, tmp: None, + payload=SUGGESTION_PAYLOAD, + deliberate_mismatch=DeliberateMismatch( + "Local git provider writes suggestions to a local artifact file and unconditionally returns True." + ), + ), + SuggestionOutcomeContract( + provider_name="plain-diff", + build_provider=_build_plain_diff_suggestion_provider, + make_fail=lambda p, mp, tmp: None, + make_succeed=lambda p, mp, tmp: None, + payload=SUGGESTION_PAYLOAD, + deliberate_mismatch=DeliberateMismatch( + "Plain-diff provider renders suggestions to stdout or output file and unconditionally returns True." + ), + ), + SuggestionOutcomeContract( + provider_name="mosaico-diff", + build_provider=lambda mp, tmp: _bare(DiffInputProvider)(mp), + make_fail=lambda p, mp, tmp: None, + make_succeed=lambda p, mp, tmp: None, + payload=SUGGESTION_PAYLOAD, + deliberate_mismatch=DeliberateMismatch("Mosaico diff is a no-op provider that unconditionally returns True."), + ), ) @@ -241,12 +511,54 @@ def _rows(tier: Tier | None = None, check_execution_only: bool = False): yield pytest.param(provider_name, contract, id=f"{provider_name}-{contract.name}") +def _predicate_rows(): + for contract in PREDICATE_CONTRACTS: + for provider_name in PROVIDERS: + yield pytest.param(provider_name, contract, id=f"{provider_name}-{contract.name}") + + +def _has_evidence(provider_type: type[GitProvider], contract: PredicateContract) -> bool: + for name in contract.evidence: + member = getattr(provider_type, name, None) + if member is not getattr(GitProvider, name, None): + return True + return False + + def test_every_registered_provider_has_a_contract_row(): contracted = {provider_type for provider_type, _ in PROVIDERS.values()} assert set(_GIT_PROVIDERS.values()) <= contracted +@pytest.mark.parametrize("provider_name,contract", tuple(_predicate_rows())) +def test_predicate_truth_matches_evidence(provider_name: str, contract: PredicateContract): + provider_type, _ = PROVIDERS[provider_name] + mismatch = contract.deliberate_mismatches.get(provider_name) + predicate = getattr(provider_type.__new__(provider_type), contract.name) + + if mismatch: + assert mismatch.reason.strip() + assert predicate() is False + else: + assert predicate() is _has_evidence(provider_type, contract) + + +@pytest.mark.parametrize("contract", PREDICATE_CONTRACTS, ids=lambda contract: contract.name) +def test_predicate_mismatches_name_only_providers_with_evidence(contract: PredicateContract): + for provider_name in contract.deliberate_mismatches: + provider_type, _ = PROVIDERS[provider_name] + + assert _has_evidence(provider_type, contract) + + +@pytest.mark.parametrize("contract", PREDICATE_CONTRACTS, ids=lambda contract: contract.name) +def test_predicate_contract_evidence_members_exist_on_registered_providers(contract: PredicateContract): + assert contract.evidence + for member_name in contract.evidence: + assert any(hasattr(provider_type, member_name) for provider_type, _ in PROVIDERS.values()) + + @pytest.mark.parametrize("contract", METHOD_CONTRACTS, ids=lambda contract: contract.name) def test_every_contract_row_places_every_provider_in_a_tier(contract: MethodContract): assert set(contract.tiers) == set(PROVIDERS) @@ -313,3 +625,62 @@ def test_disable_eyes_short_circuits_before_any_backend_call(provider_name: str) provider = provider_type.__new__(provider_type) assert provider.add_eyes_reaction(COMMENT_ID, disable_eyes=True) is None + + +@pytest.mark.parametrize( + "provider_name", + [ + name + for name, (cls, _) in PROVIDERS.items() + if "publish_code_suggestions" in cls.__dict__ + ], +) +def test_publish_code_suggestions_declares_bool_return(provider_name: str): + provider_type, _ = PROVIDERS[provider_name] + hints = get_type_hints(provider_type.publish_code_suggestions) + assert hints.get("return") is bool + + +def test_every_provider_overriding_publish_code_suggestions_has_a_contract_row(): + overriding = { + name + for name, (cls, _) in PROVIDERS.items() + if "publish_code_suggestions" in cls.__dict__ + } + contracted = {contract.provider_name for contract in SUGGESTION_OUTCOME_CONTRACTS} + assert len(contracted) == len(SUGGESTION_OUTCOME_CONTRACTS) + assert contracted == overriding + + +@pytest.mark.parametrize( + "contract", + [ + pytest.param(contract, id=f"{contract.provider_name}-publish_code_suggestions-failure") + for contract in SUGGESTION_OUTCOME_CONTRACTS + ], +) +def test_publish_code_suggestions_returns_false_on_total_failure( + contract: SuggestionOutcomeContract, monkeypatch, tmp_path +): + provider = contract.build_provider(monkeypatch, tmp_path) + contract.make_fail(provider, monkeypatch, tmp_path) + if contract.deliberate_mismatch is not None: + assert contract.deliberate_mismatch.reason.strip() + assert provider.publish_code_suggestions(contract.payload) is True + else: + assert provider.publish_code_suggestions(contract.payload) is False + + +@pytest.mark.parametrize( + "contract", + [ + pytest.param(contract, id=f"{contract.provider_name}-publish_code_suggestions-success") + for contract in SUGGESTION_OUTCOME_CONTRACTS + ], +) +def test_publish_code_suggestions_returns_true_on_success( + contract: SuggestionOutcomeContract, monkeypatch, tmp_path +): + provider = contract.build_provider(monkeypatch, tmp_path) + contract.make_succeed(provider, monkeypatch, tmp_path) + assert provider.publish_code_suggestions(contract.payload) is True diff --git a/tests/unittest/test_git_provider_utils.py b/tests/unittest/test_git_provider_utils.py index 4048f73335..6dc7105463 100644 --- a/tests/unittest/test_git_provider_utils.py +++ b/tests/unittest/test_git_provider_utils.py @@ -112,6 +112,8 @@ def test_handle_configurations_errors_uses_persistent_comment_when_supported(): assert comment["final_update_message"] is False assert "PR-Agent failed to apply 'local' repo settings" in comment["body"] assert "Invalid value" in comment["body"] + assert "https://docs.pr-agent.ai/usage-guide/configuration_options/" in comment["body"] + assert "qodo-merge-docs.qodo.ai" not in comment["body"] assert "```toml\n[config]\nmodel =\n```" in comment["body"] assert "
      Configuration content:" in comment["body"] diff --git a/tests/unittest/test_gitea_provider.py b/tests/unittest/test_gitea_provider.py index 426ae2064d..70712c9cf5 100644 --- a/tests/unittest/test_gitea_provider.py +++ b/tests/unittest/test_gitea_provider.py @@ -1189,11 +1189,16 @@ def test_gitea_persistent_wrapper_preserves_identity_and_result(): published = object() class RecordingGiteaProvider: + # Gitea no longer overrides this: the base implementation persists. publish_persistent_comment = GiteaProvider.publish_persistent_comment + supports_comment_editing = GitProvider.supports_comment_editing def __init__(self): self.calls = [] + def edit_comment(self, comment, body): + return True + def publish_persistent_comment_full( self, pr_comment, initial_header, update_header=True, name="review", final_update_message=True, as_thread=False, identity_marker=None, diff --git a/tests/unittest/test_github_app_auth.py b/tests/unittest/test_github_app_auth.py index 58e1820953..f3761d7574 100644 --- a/tests/unittest/test_github_app_auth.py +++ b/tests/unittest/test_github_app_auth.py @@ -1,3 +1,5 @@ +from unittest.mock import MagicMock + import jwt from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa @@ -8,10 +10,12 @@ class TestGithubAppAuth: def test_integer_app_id_produces_valid_jwt(self): - """GitHub App ids are integers in the settings toml, but PyJWT >=2.11 rejects a - non-string `iss` claim and PyGithub 1.59 passes the id through raw (#2955, - previously #2210; fixed upstream in PyGithub#3272, which we don't ship yet). - The provider must cast to str before building the authentication. + """GitHub App ids are integers in the settings toml. PyJWT >=2.11 rejects a + non string `iss` claim and PyGithub 1.59 passed the id through raw (#2955, + previously #2210; upstream fixed the pass through in PyGithub#3272, which + landed in 2.7.0, so the 2.10 pin ships it). The provider keeps casting to + str before building the authentication, which is thus harmless on the 2.10 + pin while it stays a hard requirement for the older 1.59 behaviour. """ key = rsa.generate_private_key(public_exponent=65537, key_size=2048) private_key_pem = key.private_bytes( @@ -46,6 +50,74 @@ def test_integer_app_id_produces_valid_jwt(self): for name, value in original.items(): settings.set(name, value) + def test_client_defaults_to_no_pacing_and_no_retry(self, monkeypatch): + """PyGithub 2.x turns on pacing and 10 retries by default; PR-Agent must + stay behaviour-neutral unless configured. The [github] settings mirror the + 1.59 defaults and the client constructor applies them explicitly + (https://github.com/The-PR-Agent/pr-agent/pull/3240).""" + settings = get_settings() + original = { + "GITHUB.DEPLOYMENT_TYPE": settings.get("GITHUB.DEPLOYMENT_TYPE", None), + "GITHUB.USER_TOKEN": settings.get("GITHUB.USER_TOKEN", None), + } + settings.set("GITHUB.DEPLOYMENT_TYPE", "user") + settings.set("GITHUB.USER_TOKEN", "stub-token") + captured = {} + + def fake_github(**kwargs): + captured.update(kwargs) + return MagicMock() + + monkeypatch.setattr("pr_agent.git_providers.github_provider.Github", fake_github) + provider = GithubProvider.__new__(GithubProvider) + provider.installation_id = None + provider.base_url = "https://api.github.com" + try: + provider._get_github_client() + finally: + for name, value in original.items(): + settings.set(name, value) + + assert captured["seconds_between_requests"] == 0 + assert captured["seconds_between_writes"] == 0 + assert captured["retry"] is None + + def test_pacing_and_retry_settings_reach_the_client(self, monkeypatch): + """Operators may opt into PyGithub 2.x throttling and backoff via the + [github] settings; each knob must land on the constructor (PR #3240).""" + settings = get_settings() + original = { + "GITHUB.DEPLOYMENT_TYPE": settings.get("GITHUB.DEPLOYMENT_TYPE", None), + "GITHUB.USER_TOKEN": settings.get("GITHUB.USER_TOKEN", None), + "GITHUB.SECONDS_BETWEEN_REQUESTS": settings.get("GITHUB.SECONDS_BETWEEN_REQUESTS", None), + "GITHUB.SECONDS_BETWEEN_WRITES": settings.get("GITHUB.SECONDS_BETWEEN_WRITES", None), + "GITHUB.API_RETRIES": settings.get("GITHUB.API_RETRIES", None), + } + settings.set("GITHUB.DEPLOYMENT_TYPE", "user") + settings.set("GITHUB.USER_TOKEN", "stub-token") + settings.set("GITHUB.SECONDS_BETWEEN_REQUESTS", 0.25) + settings.set("GITHUB.SECONDS_BETWEEN_WRITES", 1.0) + settings.set("GITHUB.API_RETRIES", 4) + captured = {} + + def fake_github(**kwargs): + captured.update(kwargs) + return MagicMock() + + monkeypatch.setattr("pr_agent.git_providers.github_provider.Github", fake_github) + provider = GithubProvider.__new__(GithubProvider) + provider.installation_id = None + provider.base_url = "https://api.github.com" + try: + provider._get_github_client() + finally: + for name, value in original.items(): + settings.set(name, value) + + assert captured["seconds_between_requests"] == 0.25 + assert captured["seconds_between_writes"] == 1.0 + assert captured["retry"].total == 4 + def test_is_bot_user_reads_github_section_setting(self): """is_bot_user must honor `ignore_bot_pr` from the [github] section (#3017): configuration.toml sets it under [github], so reading diff --git a/tests/unittest/test_github_incremental_default.py b/tests/unittest/test_github_incremental_default.py index fe707356bc..c3519c3dd4 100644 --- a/tests/unittest/test_github_incremental_default.py +++ b/tests/unittest/test_github_incremental_default.py @@ -1,4 +1,6 @@ """The default `IncrementalPR` is per call, not one object shared since import.""" +import datetime +from types import SimpleNamespace from unittest.mock import MagicMock, patch from pr_agent.git_providers.git_provider import IncrementalPR @@ -55,3 +57,101 @@ def test_an_incremental_passed_in_is_the_one_used(): provider.get_incremental_commits(given) assert provider.incremental is given + + +def test_can_run_incremental_review_handles_aware_commit_dates(): + """Regression: PyGithub 2.x returns timezone-aware commit timestamps, + but _can_run_incremental_review compared them against a naive + datetime.now(), raising TypeError on every second run.""" + from pr_agent.tools.pr_reviewer import PRReviewer + + reviewer = PRReviewer.__new__(PRReviewer) + reviewer.is_auto = False + reviewer.git_provider = MagicMock(spec=["get_incremental_commits"]) + reviewer.pr_url = "https://github.com/test/repo/pull/1" + + aware_date = datetime.datetime(2026, 9, 10, 12, 0, tzinfo=datetime.timezone.utc) + commit = SimpleNamespace(commit=SimpleNamespace(author=SimpleNamespace(date=aware_date))) + + reviewer.incremental = IncrementalPR(True) + reviewer.incremental.commits_range = [commit] + reviewer.incremental.last_seen_commit = commit + + with patch("pr_agent.tools.pr_reviewer.get_settings") as mock_settings: + mock_settings.return_value.git_provider = "github" + mock_settings.return_value.pr_reviewer.minimal_commits_for_incremental_review = 0 + mock_settings.return_value.pr_reviewer.minimal_minutes_for_incremental_review = 0 + mock_settings.return_value.pr_reviewer.require_all_thresholds_for_incremental_review = False + # Must not raise TypeError + result = reviewer._can_run_incremental_review() + + assert result in (True, False) # either outcome is valid; crash is the bug + + +def _frozen_now(fixed_aware_utc): + class _FrozenDatetime: + timezone = datetime.timezone + timedelta = datetime.timedelta + + class datetime: + @staticmethod + def now(*args, **kwargs): + return fixed_aware_utc.astimezone(datetime.timezone.utc) + + return patch("pr_agent.tools.pr_reviewer.datetime", new=_FrozenDatetime) + + +def test_incremental_review_runs_when_commit_is_older_than_utc_threshold(): + """The threshold is naive UTC, so a host in a non-UTC zone must not shift + the minimum-age decision.""" + from pr_agent.tools.pr_reviewer import PRReviewer + + reviewer = PRReviewer.__new__(PRReviewer) + reviewer.is_auto = False + reviewer.git_provider = MagicMock(spec=["get_incremental_commits"]) + reviewer.pr_url = "https://github.com/test/repo/pull/1" + + commit = SimpleNamespace( + commit=SimpleNamespace(author=SimpleNamespace(date=datetime.datetime(2026, 9, 10, 12, 0, tzinfo=datetime.timezone.utc))) + ) + reviewer.incremental = IncrementalPR(True) + reviewer.incremental.commits_range = [commit] + reviewer.incremental.last_seen_commit = commit + + with ( + _frozen_now(datetime.datetime(2026, 9, 10, 13, 0, tzinfo=datetime.timezone.utc)), + patch("pr_agent.tools.pr_reviewer.get_settings") as mock_settings, + ): + mock_settings.return_value.git_provider = "github" + mock_settings.return_value.pr_reviewer.minimal_commits_for_incremental_review = 0 + mock_settings.return_value.pr_reviewer.minimal_minutes_for_incremental_review = 30 + mock_settings.return_value.pr_reviewer.require_all_thresholds_for_incremental_review = False + + assert reviewer._can_run_incremental_review() is True + + +def test_incremental_review_skipped_when_commit_is_newer_than_utc_threshold(): + from pr_agent.tools.pr_reviewer import PRReviewer + + reviewer = PRReviewer.__new__(PRReviewer) + reviewer.is_auto = False + reviewer.git_provider = MagicMock(spec=["get_incremental_commits"]) + reviewer.pr_url = "https://github.com/test/repo/pull/1" + + commit = SimpleNamespace( + commit=SimpleNamespace(author=SimpleNamespace(date=datetime.datetime(2026, 9, 10, 13, 15, tzinfo=datetime.timezone.utc))) + ) + reviewer.incremental = IncrementalPR(True) + reviewer.incremental.commits_range = [commit] + reviewer.incremental.last_seen_commit = commit + + with ( + _frozen_now(datetime.datetime(2026, 9, 10, 13, 0, tzinfo=datetime.timezone.utc)), + patch("pr_agent.tools.pr_reviewer.get_settings") as mock_settings, + ): + mock_settings.return_value.git_provider = "github" + mock_settings.return_value.pr_reviewer.minimal_commits_for_incremental_review = 0 + mock_settings.return_value.pr_reviewer.minimal_minutes_for_incremental_review = 30 + mock_settings.return_value.pr_reviewer.require_all_thresholds_for_incremental_review = True + + assert reviewer._can_run_incremental_review() is False diff --git a/tests/unittest/test_github_issue_reference_extraction.py b/tests/unittest/test_github_issue_reference_extraction.py index c6e492125d..b0cee2db88 100644 --- a/tests/unittest/test_github_issue_reference_extraction.py +++ b/tests/unittest/test_github_issue_reference_extraction.py @@ -64,3 +64,44 @@ def test_a_full_url_is_not_bounded(): def test_a_cross_repo_shorthand_is_not_bounded(): """Control: owner/repo#123 names its repository, so it is unambiguous too.""" assert _links("Fixes other/project#12345") == [f"{BASE}/other/project/issues/12345"] + + +@pytest.mark.parametrize( + ("shorthand", "expected_url"), + [ + ("my-org/repo#42", f"{BASE}/my-org/repo/issues/42"), + ("org/my-repo#42", f"{BASE}/org/my-repo/issues/42"), + ("my-org/my-repo#42", f"{BASE}/my-org/my-repo/issues/42"), + ], +) +def test_cross_repo_with_hyphens(shorthand, expected_url): + assert _links(f"Fixes {shorthand}") == [expected_url] + + +@pytest.mark.parametrize( + ("shorthand", "expected_url"), + [ + ("org/my.repo#7", f"{BASE}/org/my.repo/issues/7"), + ("org/my_repo#7", f"{BASE}/org/my_repo/issues/7"), + ], +) +def test_cross_repo_with_period_and_underscore_in_repo_name(shorthand, expected_url): + assert _links(f"Fixes {shorthand}") == [expected_url] + + +def test_cross_repo_this_repo_own_name(): + assert _links("Fixes The-PR-Agent/pr-agent#3081") == [f"{BASE}/The-PR-Agent/pr-agent/issues/3081"] + + +def test_regression_bare_shorthand_resolves_against_current_repo(): + assert _links("Fixes #42") == [f"{BASE}/{REPO}/issues/42"] + + +def test_regression_full_url_still_wins(): + url = f"{BASE}/some-org/some-repo/issues/99" + assert _links(f"Fixes {url}") == [url] + + +def test_regression_mid_token_does_not_match_cross_repo(): + assert _links("see x/my-org/my-repo#42") == [f"{BASE}/{REPO}/issues/42"] + assert _links("see x/other/project#12345") == [f"{BASE}/{REPO}/issues/12345"] diff --git a/tests/unittest/test_github_polling.py b/tests/unittest/test_github_polling.py index ffa135ed9c..b6b9c73881 100644 --- a/tests/unittest/test_github_polling.py +++ b/tests/unittest/test_github_polling.py @@ -1,27 +1,447 @@ +import asyncio +import multiprocessing +import random +import time from collections import deque +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest from pr_agent.servers import github_polling -def test_start_queued_processes_respects_parallel_limit(monkeypatch): - created_processes = [] +@pytest.fixture +def workers(monkeypatch): + created = [] class FakeProcess: def __init__(self, target, args): self.target = target self.args = args - self.started = False - created_processes.append(self) + self.pid = None + self.alive = False + self.closed = False + self.joined = False + created.append(self) def start(self): - self.started = True + self.pid = len(created) + self.alive = True + assert sum(p.alive for p in created) <= 10 + + def is_alive(self): + assert not self.closed + return self.alive + + def join(self, timeout): + assert timeout == 0 + assert not self.alive + self.joined = True + + def close(self): + assert not self.alive + self.closed = True + + # Replace the module reference, not multiprocessing.Process process-wide. + monkeypatch.setattr(github_polling, "multiprocessing", SimpleNamespace(Process=FakeProcess)) + monkeypatch.setattr(github_polling, "get_logger", MagicMock()) + return created, FakeProcess + + +def _queue(size): + return deque((time.sleep, (i,)) for i in range(size)) + + +@pytest.mark.asyncio +async def test_start_queued_processes_respects_parallel_limit(workers): + created, _ = workers + active = [] + queue = _queue(12) + await github_polling._start_queued_processes(queue, 10, active) + assert len(created) == len(active) == 10 + assert all(process.alive for process in active) + assert [process.args for process in active] == [(i,) for i in range(10)] + assert not queue + + +@pytest.mark.asyncio +async def test_next_batch_waits_without_dropping_accepted_work(workers, monkeypatch): + created, _ = workers + active = [] + await github_polling._start_queued_processes(_queue(10), 10, active) + waiting = asyncio.Event() + release = asyncio.Event() + + async def wait_for_capacity(delay): + assert delay == 0.25 + waiting.set() + await release.wait() + + monkeypatch.setattr(github_polling, "asyncio", SimpleNamespace(sleep=wait_for_capacity)) + queue = _queue(3) + task = asyncio.create_task(github_polling._start_queued_processes(queue, 10, active)) + try: + await asyncio.wait_for(waiting.wait(), 2) + assert len(created) == 10 + assert len(queue) == 3 + for process in created[:3]: + process.alive = False + release.set() + await asyncio.wait_for(task, 2) + assert len(created) == 13 + assert len(active) == 10 + assert all(p.joined and p.closed for p in created[:3]) + assert not queue + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_cancelled_capacity_wait_does_not_start_or_discard_work(workers, monkeypatch): + created, _ = workers + active = [] + await github_polling._start_queued_processes(_queue(10), 10, active) + entered = asyncio.Event() + + async def wait_for_capacity(delay): + entered.set() + await asyncio.Event().wait() + + monkeypatch.setattr(github_polling, "asyncio", SimpleNamespace(sleep=wait_for_capacity)) + queue = _queue(2) + task = asyncio.create_task(github_polling._start_queued_processes(queue, 10, active)) + try: + await asyncio.wait_for(entered.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await task + assert len(created) == len(active) == 10 + assert len(queue) == 2 + github_polling.get_logger().error.assert_called_once_with( + "Polling dispatch stopped with 2 tasks not dispatched" + ) + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("partial_start", [False, True]) +async def test_start_failure_keeps_previous_workers_tracked(workers, monkeypatch, partial_start): + created, process_type = workers + start = process_type.start + + def fail_second(process): + if len(created) == 2: + if partial_start: + start(process) + raise OSError("Cannot start worker") + start(process) + + monkeypatch.setattr(process_type, "start", fail_second) + active = [] + queue = _queue(3) + with pytest.raises(github_polling._PollingWorkerStartError, match="startup failed") as failure: + await github_polling._start_queued_processes(queue, 10, active) + assert isinstance(failure.value.__cause__, OSError) + assert active == (created if partial_start else created[:1]) + assert created[0].alive + assert created[1].closed is (not partial_start) + expected_remaining = 1 if partial_start else 2 + assert len(queue) == expected_remaining + github_polling.get_logger().error.assert_called_once_with( + f"Polling dispatch stopped with {expected_remaining} tasks not dispatched" + ) + assert len(created) == 2 + + +@pytest.mark.asyncio +async def test_many_batches_keep_the_active_limit(workers, monkeypatch): + created, _ = workers + active = [] + rng = random.Random(42) + accepted = 0 + + async def complete_workers(delay): + for process in active[:3]: + process.alive = False + + monkeypatch.setattr(github_polling, "asyncio", SimpleNamespace(sleep=complete_workers)) + for _ in range(100): + size = rng.randint(1, 12) + accepted += min(size, 10) + await github_polling._start_queued_processes(_queue(size), 10, active) + assert len(active) <= 10 + for process in active[:rng.randint(0, len(active))]: + process.alive = False + assert len(created) == accepted + for process in active: + process.alive = False + github_polling._reap_finished_processes(active) + assert not active + assert all(p.joined and p.closed for p in created) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("idle_response", [304, 200, 500, OSError("poll failed")]) +async def test_polling_loop_reaps_workers_on_idle_and_failed_polls(workers, monkeypatch, idle_response): + created, _ = workers + settings = SimpleNamespace( + github=SimpleNamespace(deployment_type="user", user_token="test-token"), set=MagicMock() + ) + monkeypatch.setattr(github_polling, "get_settings", lambda: settings) + monkeypatch.setattr( + github_polling, + "get_git_provider", + lambda: lambda: SimpleNamespace(get_user_id=lambda: "bot"), + ) + monkeypatch.setattr(github_polling, "mark_notification_as_read", AsyncMock()) + monkeypatch.setattr(github_polling, "is_valid_notification", AsyncMock(return_value=( + True, set(), {"id": 2}, "@bot /review", "https://example.test/pull/1", "@bot" + ))) + finished = asyncio.Event() + polls = 0 + sleeps = 0 + + class Response: + def __init__(self, status): + self.status = status + self.headers = {} + + async def __aenter__(self): + if isinstance(self.status, Exception): + raise self.status + return self + + async def __aexit__(self, *args): + pass + + async def json(self): + return [{"id": 1}] if polls == 1 else [] + + class Session: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + def get(self, *args, **kwargs): + nonlocal polls + polls += 1 + return Response(200 if polls == 1 else idle_response) + + async def sleep(delay): + nonlocal sleeps + sleeps += 1 + if sleeps == 2: + created[0].alive = False + if sleeps == 3: + finished.set() + await asyncio.Event().wait() + + monkeypatch.setattr(github_polling, "aiohttp", SimpleNamespace(ClientSession=Session)) + monkeypatch.setattr(github_polling, "asyncio", SimpleNamespace(sleep=sleep)) + task = asyncio.create_task(github_polling.polling_loop()) + try: + await asyncio.wait_for(finished.wait(), 2) + assert len(created) == 1 + assert created[0].joined and created[0].closed + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_reap_real_spawned_worker(monkeypatch): + ctx = multiprocessing.get_context("spawn") + monkeypatch.setattr(github_polling, "multiprocessing", SimpleNamespace(Process=ctx.Process)) + active = [] + await github_polling._start_queued_processes(deque([(time.sleep, (0,))]), 1, active) + process = active[0] + try: + await asyncio.to_thread(process.join, 5) + assert not process.is_alive() + github_polling._reap_finished_processes(active) + assert not active + with pytest.raises(ValueError, match="closed"): + process.is_alive() + finally: + if active: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + process.close() + + +@pytest.mark.asyncio +async def test_polling_loop_shares_capacity_between_batches(workers, monkeypatch): + created, _ = workers + settings = SimpleNamespace( + github=SimpleNamespace(deployment_type="user", user_token="test-token"), set=MagicMock() + ) + monkeypatch.setattr(github_polling, "get_settings", lambda: settings) + monkeypatch.setattr( + github_polling, + "get_git_provider", + lambda: lambda: SimpleNamespace(get_user_id=lambda: "bot"), + ) + monkeypatch.setattr(github_polling, "mark_notification_as_read", AsyncMock()) + monkeypatch.setattr(github_polling, "is_valid_notification", AsyncMock(return_value=( + True, set(), {"id": 2}, "@bot /review", "https://example.test/pull/1", "@bot" + ))) + waiting = asyncio.Event() + release = asyncio.Event() + finished = asyncio.Event() + polls = 0 + + class Session: + status = 200 + headers = {} + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + def get(self, *args, **kwargs): + nonlocal polls + polls += 1 + return self + + async def json(self): + return [{"id": i} for i in range(10 if polls == 1 else 2)] + + async def sleep(delay): + if delay == 0.25: + waiting.set() + await release.wait() + elif polls == 2: + finished.set() + await asyncio.Event().wait() + + monkeypatch.setattr(github_polling, "aiohttp", SimpleNamespace(ClientSession=Session)) + monkeypatch.setattr(github_polling, "asyncio", SimpleNamespace(sleep=sleep)) + task = asyncio.create_task(github_polling.polling_loop()) + try: + await asyncio.wait_for(waiting.wait(), 2) + assert len(created) == 10 + assert polls == 2 + for process in created[:2]: + process.alive = False + release.set() + await asyncio.wait_for(finished.wait(), 2) + assert len(created) == 12 + assert sum(process.alive for process in created) == 10 + finally: + task.cancel() + result, = await asyncio.gather(task, return_exceptions=True) + assert isinstance(result, asyncio.CancelledError) + + +@pytest.mark.asyncio +async def test_polling_loop_logs_accepted_work_cancelled_before_dispatch(workers, monkeypatch): + created, _ = workers + settings = SimpleNamespace( + github=SimpleNamespace(deployment_type="user", user_token="test-token"), set=MagicMock() + ) + monkeypatch.setattr(github_polling, "get_settings", lambda: settings) + monkeypatch.setattr( + github_polling, + "get_git_provider", + lambda: lambda: SimpleNamespace(get_user_id=lambda: "bot"), + ) + monkeypatch.setattr(github_polling, "mark_notification_as_read", AsyncMock()) + entered_second = asyncio.Event() + validation_calls = 0 + + async def validate(*args, **kwargs): + nonlocal validation_calls + validation_calls += 1 + if validation_calls == 1: + return True, set(), {"id": 2}, "@bot /review", "https://example.test/pull/1", "@bot" + entered_second.set() + await asyncio.Event().wait() + raise AssertionError("Second validation should remain blocked") + + monkeypatch.setattr(github_polling, "is_valid_notification", validate) + + class Session: + status = 200 + headers = {} + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + def get(self, *args, **kwargs): + return self + + async def json(self): + return [{"id": 1}, {"id": 2}] + + async def sleep(delay): + await asyncio.sleep(0) + + monkeypatch.setattr(github_polling, "aiohttp", SimpleNamespace(ClientSession=Session)) + monkeypatch.setattr(github_polling, "asyncio", SimpleNamespace(sleep=sleep)) + task = asyncio.create_task(github_polling.polling_loop()) + await asyncio.wait_for(entered_second.wait(), 2) + task.cancel() + result, = await asyncio.gather(task, return_exceptions=True) + + assert isinstance(result, asyncio.CancelledError) + assert not created + github_polling.get_logger().error.assert_any_call( + "Polling dispatch stopped with 1 tasks not dispatched" + ) + + +@pytest.mark.asyncio +async def test_polling_loop_stops_after_worker_start_failure(workers, monkeypatch): + _, process_type = workers + settings = SimpleNamespace( + github=SimpleNamespace(deployment_type="user", user_token="test-token"), set=MagicMock() + ) + monkeypatch.setattr(github_polling, "get_settings", lambda: settings) + monkeypatch.setattr( + github_polling, + "get_git_provider", + lambda: lambda: SimpleNamespace(get_user_id=lambda: "bot"), + ) + monkeypatch.setattr(github_polling, "mark_notification_as_read", AsyncMock()) + monkeypatch.setattr(github_polling, "is_valid_notification", AsyncMock(return_value=( + True, set(), {"id": 2}, "@bot /review", "https://example.test/pull/1", "@bot" + ))) + start = MagicMock(side_effect=OSError("Cannot start worker")) + monkeypatch.setattr(process_type, "start", start) + + class Session: + status = 200 + headers = {} + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + def get(self, *args, **kwargs): + return self - monkeypatch.setattr(github_polling.multiprocessing, "Process", FakeProcess) + async def json(self): + return [{"id": 1}] - task_queue = deque((lambda: None, ()) for _ in range(12)) - processes = github_polling._start_queued_processes(task_queue, 10) + async def sleep(delay): + await asyncio.sleep(0) - assert len(created_processes) == 10 - assert len(processes) == 10 - assert all(process.started for process in processes) - assert not task_queue + monkeypatch.setattr(github_polling, "aiohttp", SimpleNamespace(ClientSession=Session)) + monkeypatch.setattr(github_polling, "asyncio", SimpleNamespace(sleep=sleep)) + with pytest.raises(github_polling._PollingWorkerStartError, match="startup failed"): + await asyncio.wait_for(github_polling.polling_loop(), timeout=2) + start.assert_called_once() diff --git a/tests/unittest/test_github_provider_comments.py b/tests/unittest/test_github_provider_comments.py index 8b686acb58..599704c694 100644 --- a/tests/unittest/test_github_provider_comments.py +++ b/tests/unittest/test_github_provider_comments.py @@ -55,6 +55,30 @@ def _make_provider(pr=None, max_chars=65000): return p +@pytest.mark.parametrize( + ("output", "max_chars", "expected"), + [ + pytest.param("short", 10, "short", id="short"), + pytest.param("exact", 5, "exact", id="exact"), + pytest.param("x" * 20, 10, ("x" * 7) + "...", id="truncated"), + pytest.param("abcdef", -1, "", id="negative-limit"), + pytest.param("abcdef", 0, "", id="zero-limit"), + pytest.param("abcdef", 1, ".", id="one-character-limit"), + pytest.param("abcdef", 2, "..", id="two-character-limit"), + pytest.param("abcdef", 3, "...", id="three-character-limit"), + pytest.param("😀" * 5, 4, "😀...", id="unicode-code-points"), + ], +) +def test_limit_output_characters_respects_total_limit(output, max_chars, expected): + provider = _make_provider() + + result = provider.limit_output_characters(output, max_chars) + + assert result == expected + if max_chars >= 0: + assert len(result) <= max_chars + + def test_edit_comment_returns_false_on_github_failure(): provider = _make_provider() comment = MagicMock() @@ -172,8 +196,8 @@ def test_create_inline_comment_limits_body_length(monkeypatch): payload = provider.create_inline_comment(long_body, "f.py", "line") assert payload["body"].endswith("...") - # limit_output_characters: output[:max_chars] + '...' - assert payload["body"] == "A" * 10 + "..." + assert payload["body"] == "A" * 7 + "..." + assert len(payload["body"]) == provider.max_comment_chars def test_create_inline_comment_does_not_truncate_short_body(monkeypatch): @@ -312,6 +336,7 @@ def test_publish_code_suggestions_multi_line_payload_shape(): def capture(comments, disable_fallback=False): captured["comments"] = comments + return True provider.publish_inline_comments = capture @@ -416,6 +441,209 @@ def boom(comments, disable_fallback=False): assert result is False +def test_publish_code_suggestions_422_fallback_all_dropped_returns_false(monkeypatch): + """Regression test for #3223: When the 422 fallback drops all comments + (0 verified, 0 repaired), publish_code_suggestions must return False so caller + can trigger retry logic.""" + fake_pr = _FakePR(raise_on_first=_FakeGithubException(status=422)) + provider = _make_provider(pr=fake_pr) + _stub_validation_passthrough(provider) + + # All comments are rejected during verification + monkeypatch.setattr( + provider, + "_verify_code_comments", + lambda comments: ([], [(c, Exception("invalid")) for c in comments]), + ) + # No invalid comment can be repaired + monkeypatch.setattr( + provider, + "_try_fix_invalid_inline_comments", + lambda invalid_list: [], + ) + + suggestions = [{ + "body": "```suggestion\nsuggestion\n```", + "relevant_file": "src/foo.py", + "relevant_lines_start": 10, + "relevant_lines_end": 12, + }] + + result = provider.publish_code_suggestions(suggestions) + assert result is False + # Only the initial failing create_review call occurred; 0 fallback comments posted + assert len(fake_pr.create_review_calls) == 1 + + +def test_publish_code_suggestions_422_fallback_partial_success_returns_true(monkeypatch): + """When 422 fallback successfully publishes at least one comment, return True + to prevent duplicate comment creation by whole-batch retries.""" + fake_pr = _FakePR(raise_on_first=_FakeGithubException(status=422)) + provider = _make_provider(pr=fake_pr) + _stub_validation_passthrough(provider) + + # 1 verified comment, 1 invalid comment + def fake_verify(comments): + return [comments[0]], [(comments[1], Exception("invalid"))] + + monkeypatch.setattr(provider, "_verify_code_comments", fake_verify) + monkeypatch.setattr(provider, "_try_fix_invalid_inline_comments", lambda invalid_list: []) + + suggestions = [ + { + "body": "```suggestion\nfirst\n```", + "relevant_file": "src/foo.py", + "relevant_lines_start": 1, + "relevant_lines_end": 2, + }, + { + "body": "```suggestion\nsecond\n```", + "relevant_file": "src/foo.py", + "relevant_lines_start": 5, + "relevant_lines_end": 6, + }, + ] + + result = provider.publish_code_suggestions(suggestions) + assert result is True + # 1 initial failed batch call, 1 successful fallback call with the verified comment + assert len(fake_pr.create_review_calls) == 2 + assert len(fake_pr.create_review_calls[1]["comments"]) == 1 + + +def test_publish_code_suggestions_422_fallback_repaired_comment_success(monkeypatch): + """When initial batch gets 422, verification rejects, but repairing succeeds and + individual publish succeeds, return True.""" + fake_pr = _FakePR(raise_on_first=_FakeGithubException(status=422)) + provider = _make_provider(pr=fake_pr) + _stub_validation_passthrough(provider) + + settings = SimpleNamespace( + github=SimpleNamespace(try_fix_invalid_inline_comments=True), + get=lambda key, default=None: default, + ) + monkeypatch.setattr(gh_module, "get_settings", lambda: settings) + + monkeypatch.setattr( + provider, + "_verify_code_comments", + lambda comments: ([], [(comments[0], Exception("invalid"))]), + ) + repaired = [{"body": "fixed single line", "path": "src/foo.py", "line": 10, "side": "RIGHT"}] + monkeypatch.setattr(provider, "_try_fix_invalid_inline_comments", lambda invalid: repaired) + + suggestions = [{ + "body": "```suggestion\nmulti\nline\n```", + "relevant_file": "src/foo.py", + "relevant_lines_start": 10, + "relevant_lines_end": 12, + }] + + result = provider.publish_code_suggestions(suggestions) + assert result is True + # Call 1: initial batch -> raises 422 + # Call 2: repaired comment via publish_inline_comments([comment], disable_fallback=True) -> succeeds + assert len(fake_pr.create_review_calls) == 2 + assert fake_pr.create_review_calls[1]["comments"] == repaired + + +def test_publish_code_suggestions_422_fallback_repaired_comment_failure_returns_false(monkeypatch): + """When repaired payload is generated but publishing that repaired comment fails, + it must NOT count as published, and publish_code_suggestions must return False.""" + fake_pr = _FakePR(raise_on_first=_FakeGithubException(status=422)) + provider = _make_provider(pr=fake_pr) + _stub_validation_passthrough(provider) + + settings = SimpleNamespace( + github=SimpleNamespace(try_fix_invalid_inline_comments=True), + get=lambda key, default=None: default, + ) + monkeypatch.setattr(gh_module, "get_settings", lambda: settings) + + monkeypatch.setattr( + provider, + "_verify_code_comments", + lambda comments: ([], [(comments[0], Exception("invalid"))]), + ) + repaired = [{"body": "fixed single line", "path": "src/foo.py", "line": 10, "side": "RIGHT"}] + monkeypatch.setattr(provider, "_try_fix_invalid_inline_comments", lambda invalid: repaired) + + calls = 0 + + def fail_repaired(commit=None, comments=None): + nonlocal calls + calls += 1 + if calls == 1: + raise _FakeGithubException(status=422) + # with disable_fallback=True, this re-raises from publish_inline_comments + raise _FakeGithubException(status=422) + + fake_pr.create_review = fail_repaired + + suggestions = [{ + "body": "```suggestion\nmulti\nline\n```", + "relevant_file": "src/foo.py", + "relevant_lines_start": 10, + "relevant_lines_end": 12, + }] + + result = provider.publish_code_suggestions(suggestions) + assert result is False + assert calls == 2 + + +def test_publish_code_suggestions_normal_success(): + """Clean create_review call without 422 must return True.""" + fake_pr = _FakePR() + provider = _make_provider(pr=fake_pr) + _stub_validation_passthrough(provider) + + suggestions = [{ + "body": "normal", + "relevant_file": "src/foo.py", + "relevant_lines_start": 1, + "relevant_lines_end": 1, + }] + + result = provider.publish_code_suggestions(suggestions) + assert result is True + assert len(fake_pr.create_review_calls) == 1 + + +def test_persistent_dedup_all_skipped_returns_true(monkeypatch): + """When persistent_inline_comments is enabled and all comments are duplicates, + publish_inline_comments and publish_code_suggestions must return True without + calling create_review.""" + fake_pr = _FakePR() + provider = _make_provider(pr=fake_pr) + _stub_validation_passthrough(provider) + + settings = SimpleNamespace( + get=lambda key, default=None: True if key == "config.persistent_inline_comments" else default, + github=SimpleNamespace(try_fix_invalid_inline_comments=False), + ) + monkeypatch.setattr(gh_module, "get_settings", lambda: settings) + + store = MagicMock() + store.seen.return_value = True + monkeypatch.setattr(gh_module, "get_inline_comment_store", lambda prov: store) + + comments = [{"path": "src/foo.py", "body": "already posted", "line": 5}] + res_inline = provider.publish_inline_comments(comments) + assert res_inline is True + assert len(fake_pr.create_review_calls) == 0 + + suggestions = [{ + "body": "already posted", + "relevant_file": "src/foo.py", + "relevant_lines_start": 5, + "relevant_lines_end": 5, + }] + res_suggestions = provider.publish_code_suggestions(suggestions) + assert res_suggestions is True + assert len(fake_pr.create_review_calls) == 0 + + # --------------------------------------------------------------------------- # resolve_comment_thread # --------------------------------------------------------------------------- diff --git a/tests/unittest/test_github_webhook_delivery.py b/tests/unittest/test_github_webhook_delivery.py new file mode 100644 index 0000000000..5826e30ebe --- /dev/null +++ b/tests/unittest/test_github_webhook_delivery.py @@ -0,0 +1,230 @@ +"""Tests for opt-in GitHub webhook delivery deduplication.""" + +import asyncio +import copy +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from starlette.background import BackgroundTasks +from starlette_context import request_cycle_context + +from pr_agent.config_loader import global_settings +from pr_agent.servers import github_app +from pr_agent.servers import utils as servers_utils + + +@pytest.fixture +def delivery_clock(monkeypatch): + now = [0.0] + monkeypatch.setattr( + servers_utils.DefaultDictWithTimeout, + "_DefaultDictWithTimeout__time", + staticmethod(lambda: now[0]), + ) + monkeypatch.setattr(github_app, "time", SimpleNamespace(monotonic=lambda: now[0])) + return now + + +@pytest.fixture +def delivery_context(monkeypatch, delivery_clock): + settings = copy.deepcopy(global_settings) + settings.set("GITHUB_APP.WEBHOOK_DELIVERY_DEDUPLICATION", True) + settings.set("GITHUB_APP.PUSH_TRIGGER_PENDING_TASKS_TTL", 300) + monkeypatch.setattr(github_app, "_WEBHOOK_DELIVERY_TTL", 300) + monkeypatch.setattr(servers_utils, "_push_trigger_states_by_ttl", {}) + monkeypatch.setattr(servers_utils, "_active_push_trigger_states", {}) + monkeypatch.setattr( + github_app, + "_completed_webhook_deliveries", + servers_utils.DefaultDictWithTimeout( + float, ttl=300, update_key_time_on_get=False + ), + ) + with request_cycle_context({"settings": settings}): + yield settings + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delivery_id", ["delivery-1", None]) +async def test_github_webhook_route_forwards_delivery_id(monkeypatch, delivery_id): + body = {"installation": {"id": 1}, "action": "created"} + observed = [] + + class Request: + headers = {"X-GitHub-Event": "issue_comment"} + if delivery_id is not None: + headers["X-GitHub-Delivery"] = delivery_id + + async def fake_get_body(_request): + return body + + async def fake_handle_request(*args, **kwargs): + observed.append((args, kwargs)) + + monkeypatch.setattr(github_app, "get_body", fake_get_body) + monkeypatch.setattr(github_app, "handle_request", fake_handle_request) + background_tasks = BackgroundTasks() + + with request_cycle_context({}): + await github_app.handle_github_webhooks(background_tasks, Request(), object()) + await background_tasks() + + assert observed == [ + ((body,), {"event": "issue_comment", "delivery_id": delivery_id}), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("enabled", "delivery_id"), [(False, "delivery-1"), (True, None), (True, "")]) +async def test_handle_request_bypasses_deduplication_when_disabled_or_missing_id( + monkeypatch, delivery_context, enabled, delivery_id +): + delivery_context.set("GITHUB_APP.WEBHOOK_DELIVERY_DEDUPLICATION", enabled) + dispatch = AsyncMock(return_value={}) + monkeypatch.setattr(github_app, "_dispatch_request", dispatch) + body = {"action": "created"} + + assert await github_app.handle_request(body, "issue_comment", delivery_id=delivery_id) == {} + assert await github_app.handle_request(body, "issue_comment", delivery_id=delivery_id) == {} + + assert dispatch.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("result", [True, {}, None]) +async def test_handle_request_suppresses_completed_or_intentionally_ignored_delivery( + monkeypatch, delivery_context, result +): + dispatch = AsyncMock(return_value=result) + monkeypatch.setattr(github_app, "_dispatch_request", dispatch) + body = {"action": "created"} + + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + + dispatch.assert_awaited_once_with(body, "issue_comment", "created") + + +@pytest.mark.asyncio +async def test_handle_request_keeps_distinct_delivery_ids_independent(monkeypatch, delivery_context): + dispatch = AsyncMock(return_value={}) + monkeypatch.setattr(github_app, "_dispatch_request", dispatch) + body = {"action": "created"} + + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-2") + + assert dispatch.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("elapsed", [0, 301]) +async def test_handle_request_suppresses_active_delivery_even_after_ttl( + monkeypatch, delivery_context, delivery_clock, elapsed +): + entered = asyncio.Event() + release = asyncio.Event() + + async def dispatch(_body, _event, _action): + entered.set() + await release.wait() + return {} + + dispatch_mock = AsyncMock(side_effect=dispatch) + monkeypatch.setattr(github_app, "_dispatch_request", dispatch_mock) + body = {"action": "created"} + first = asyncio.create_task(github_app.handle_request(body, "issue_comment", delivery_id="delivery-1")) + try: + await asyncio.wait_for(entered.wait(), 1) + delivery_clock[0] += elapsed + await asyncio.wait_for( + github_app.handle_request(body, "issue_comment", delivery_id="delivery-1"), 1 + ) + dispatch_mock.assert_awaited_once() + finally: + release.set() + await asyncio.wait_for(first, 1) + + assert not servers_utils._active_push_trigger_states + + +@pytest.mark.asyncio +async def test_active_delivery_does_not_block_a_different_delivery(monkeypatch, delivery_context): + entered = asyncio.Event() + release = asyncio.Event() + + async def dispatch(body, _event, _action): + if body["comment_id"] == 1: + entered.set() + await release.wait() + return {} + + dispatch_mock = AsyncMock(side_effect=dispatch) + monkeypatch.setattr(github_app, "_dispatch_request", dispatch_mock) + first = asyncio.create_task( + github_app.handle_request({"action": "created", "comment_id": 1}, "issue_comment", delivery_id="delivery-1") + ) + try: + await asyncio.wait_for(entered.wait(), 1) + await asyncio.wait_for( + github_app.handle_request({"action": "created", "comment_id": 2}, "issue_comment", delivery_id="delivery-2"), + 1, + ) + assert dispatch_mock.await_count == 2 + finally: + release.set() + await asyncio.wait_for(first, 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [RuntimeError("agent failed"), asyncio.CancelledError()]) +async def test_failed_or_cancelled_dispatch_can_retry(monkeypatch, delivery_context, error): + dispatch = AsyncMock(side_effect=[error, {}]) + monkeypatch.setattr(github_app, "_dispatch_request", dispatch) + body = {"action": "created"} + + with pytest.raises(type(error)): + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + assert not servers_utils._active_push_trigger_states + + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + + assert dispatch.await_count == 2 + + +@pytest.mark.asyncio +async def test_duplicate_does_not_extend_completed_delivery_ttl(monkeypatch, delivery_context, delivery_clock): + dispatch = AsyncMock(return_value={}) + monkeypatch.setattr(github_app, "_dispatch_request", dispatch) + body = {"action": "created"} + + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + delivery_clock[0] = 299 + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + assert dispatch.await_count == 1 + + delivery_clock[0] = 301 + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + assert dispatch.await_count == 2 + + +@pytest.mark.asyncio +async def test_completed_delivery_ttl_starts_after_dispatch(monkeypatch, delivery_context, delivery_clock): + async def dispatch(_body, _event, _action): + delivery_clock[0] += 301 + return {} + + dispatch_mock = AsyncMock(side_effect=dispatch) + monkeypatch.setattr(github_app, "_dispatch_request", dispatch_mock) + body = {"action": "created"} + + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + delivery_clock[0] = 600 + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + assert dispatch_mock.await_count == 1 + + delivery_clock[0] = 602 + await github_app.handle_request(body, "issue_comment", delivery_id="delivery-1") + assert dispatch_mock.await_count == 2 diff --git a/tests/unittest/test_github_webhook_delivery_failures.py b/tests/unittest/test_github_webhook_delivery_failures.py new file mode 100644 index 0000000000..6df03c21ca --- /dev/null +++ b/tests/unittest/test_github_webhook_delivery_failures.py @@ -0,0 +1,142 @@ +"""Exercise failed delivery retries through the real GitHub event handlers.""" + +import asyncio +import copy +from unittest.mock import AsyncMock, Mock, call + +import pytest +from starlette_context import request_cycle_context + +from pr_agent.config_loader import get_settings +from pr_agent.identity_providers.identity_provider import Eligibility +from pr_agent.servers import github_app +from pr_agent.servers import utils as servers_utils + +EVENTS = [ + ("issue_comment", "created"), + ("pull_request", "opened"), + ("pull_request_review", "submitted"), + ("pull_request", "synchronize"), +] +AUTO_EVENTS = EVENTS[1:] +PR_URL = "https://api.github.com/repos/org/repo/pulls/1" + + +@pytest.fixture +def delivery_agent(monkeypatch): + settings = copy.deepcopy(get_settings()) + settings.set("GITHUB_APP.WEBHOOK_DELIVERY_DEDUPLICATION", True) + settings.set("GITHUB_APP.HANDLE_PR_ACTIONS", ["opened"]) + settings.set("GITHUB_APP.HANDLE_PUSH_TRIGGER", True) + settings.set("GITHUB_APP.PUSH_TRIGGER_IGNORE_MERGE_COMMITS", False) + settings.set("GITHUB_APP.FEEDBACK_ON_DRAFT_PR", False) + settings.set("GITHUB_APP.REVIEW_AUTHOR_TYPES", ["User"]) + settings.set("GITHUB_APP.REVIEW_STATES", ["changes_requested"]) + settings.set("GITHUB_APP.PR_COMMANDS", ["/review"]) + settings.set("GITHUB_APP.REVIEW_COMMANDS", ["/review"]) + settings.set("GITHUB_APP.PUSH_COMMANDS", ["/review"]) + settings.set("CONFIG.DISABLE_AUTO_FEEDBACK", False) + agent = Mock(handle_request=AsyncMock(return_value=True)) + identity = Mock() + identity.verify_eligibility.return_value = Eligibility.ELIGIBLE + monkeypatch.setattr(github_app, "PRAgent", lambda: agent) + monkeypatch.setattr(github_app, "get_identity_provider", lambda: identity) + monkeypatch.setattr(github_app, "get_git_provider_with_context", lambda **kwargs: Mock()) + monkeypatch.setattr(github_app, "apply_repo_settings", lambda api_url: None) + monkeypatch.setattr(servers_utils, "_push_trigger_states_by_ttl", {}) + monkeypatch.setattr(servers_utils, "_active_push_trigger_states", {}) + monkeypatch.setattr( + github_app, + "_completed_webhook_deliveries", + servers_utils.DefaultDictWithTimeout( + float, ttl=github_app._WEBHOOK_DELIVERY_TTL, update_key_time_on_get=False + ), + ) + with request_cycle_context({"settings": settings}): + yield agent, "delivery-1" + + +def _body(action): + return { + "action": action, + "installation": {"id": 1}, + "sender": {"id": 2, "login": "user", "type": "User"}, + "repository": {"full_name": "org/repo"}, + "pull_request": { + "url": PR_URL, + "state": "open", + "draft": False, + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-02T00:00:00Z", + }, + "comment": {"body": "/review", "pull_request_url": PR_URL, "id": 42}, + "review": {"state": "changes_requested", "user": {"type": "User"}}, + "before": "before-sha", + "after": "after-sha", + } + + +@pytest.mark.parametrize(("event", "action"), EVENTS) +async def test_false_agent_result_can_be_retried(delivery_agent, event, action): + agent, delivery_id = delivery_agent + agent.handle_request.side_effect = [False, True] + body = _body(action) + + await github_app.handle_request(body, event, delivery_id) + await github_app.handle_request(body, event, delivery_id) + await github_app.handle_request(body, event, delivery_id) + + assert agent.handle_request.await_count == 2 + + +@pytest.mark.parametrize(("event", "action"), AUTO_EVENTS) +@pytest.mark.parametrize("failure", [False, RuntimeError("agent failed")], ids=["false", "exception"]) +async def test_auto_commands_continue_after_failure_and_allow_retry(delivery_agent, event, action, failure): + agent, delivery_id = delivery_agent + for commands in ("PR_COMMANDS", "REVIEW_COMMANDS", "PUSH_COMMANDS"): + get_settings().set(f"GITHUB_APP.{commands}", ["/review", "/describe"]) + agent.handle_request.side_effect = [failure, True, True, True] + body = _body(action) + + await github_app.handle_request(body, event, delivery_id) + await github_app.handle_request(body, event, delivery_id) + await github_app.handle_request(body, event, delivery_id) + + assert agent.handle_request.await_args_list == [ + call(PR_URL, ["/review"]), + call(PR_URL, ["/describe"]), + call(PR_URL, ["/review"]), + call(PR_URL, ["/describe"]), + ] + + +@pytest.mark.parametrize(("event", "action"), EVENTS) +async def test_cancelled_agent_call_propagates_and_allows_retry(delivery_agent, event, action): + agent, delivery_id = delivery_agent + agent.handle_request.side_effect = [asyncio.CancelledError(), True] + body = _body(action) + + with pytest.raises(asyncio.CancelledError): + await github_app.handle_request(body, event, delivery_id) + await github_app.handle_request(body, event, delivery_id) + await github_app.handle_request(body, event, delivery_id) + + assert agent.handle_request.await_count == 2 + + +@pytest.mark.parametrize(("event", "action"), EVENTS) +async def test_intentionally_ignored_delivery_is_completed(delivery_agent, monkeypatch, event, action): + agent, delivery_id = delivery_agent + body = _body(action) + if event == "issue_comment": + body["comment"]["body"] = "Thanks for the review" + else: + body["pull_request"]["draft"] = True + dispatch = AsyncMock(wraps=github_app._dispatch_request) + monkeypatch.setattr(github_app, "_dispatch_request", dispatch) + + await github_app.handle_request(body, event, delivery_id) + await github_app.handle_request(body, event, delivery_id) + + dispatch.assert_awaited_once() + agent.handle_request.assert_not_awaited() diff --git a/tests/unittest/test_gitlab_file_read_efficiency.py b/tests/unittest/test_gitlab_file_read_efficiency.py new file mode 100644 index 0000000000..6353c61331 --- /dev/null +++ b/tests/unittest/test_gitlab_file_read_efficiency.py @@ -0,0 +1,22 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +from pr_agent.git_providers.gitlab_provider import GitLabProvider + + +def test_get_pr_file_content_uses_lazy_project_handle(): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/repo" + + file_obj = MagicMock() + file_obj.decode.return_value = b"hello\n" + project = MagicMock() + project.files.get.return_value = file_obj + projects = MagicMock() + projects.get.return_value = project + provider.gl = SimpleNamespace(projects=projects) + + assert provider.get_pr_file_content("src/app.py", "deadbeef") == "hello\n" + + projects.get.assert_called_once_with("group/repo", lazy=True) + project.files.get.assert_called_once_with("src/app.py", "deadbeef") diff --git a/tests/unittest/test_gitlab_provider.py b/tests/unittest/test_gitlab_provider.py index 1e06cd786c..98d6b83c37 100644 --- a/tests/unittest/test_gitlab_provider.py +++ b/tests/unittest/test_gitlab_provider.py @@ -1929,7 +1929,6 @@ def test_persistent_state_ownership_uses_authenticated_user_id(self): @pytest.mark.parametrize("capability", [ "create_inline_comment", "publish_inline_comments", - "publish_file_comments", ]) def test_unimplemented_capabilities_stay_unsupported(self, capability): assert self._provider().is_supported(capability) is False diff --git a/tests/unittest/test_gitlab_ticket_requests.py b/tests/unittest/test_gitlab_ticket_requests.py new file mode 100644 index 0000000000..37389916f1 --- /dev/null +++ b/tests/unittest/test_gitlab_ticket_requests.py @@ -0,0 +1,125 @@ +"""Exercise ticket extraction through python-gitlab's real HTTP request construction.""" +import asyncio +import json +from types import SimpleNamespace +from urllib.parse import urlparse + +import gitlab +import pytest +import requests +from requests.adapters import BaseAdapter + +from pr_agent.tools import ticket_pr_compliance_check as tickets + + +class IssueTransport(BaseAdapter): + def __init__(self, outcomes=()): + super().__init__() + self.requests = [] + self.outcomes = iter(outcomes) + + def send(self, request, **kwargs): + self.requests.append(request) + outcome = next(self.outcomes, 200) + if isinstance(outcome, BaseException): + raise outcome + # A project metadata request must never reach the transport. + assert request.method == "GET" and "/issues/" in request.url + iid = int(urlparse(request.url).path.rsplit("/", 1)[-1]) + payload = { + "id": iid + 100, + "iid": iid, + "web_url": f"https://gitlab.example/group/repo/-/issues/{iid}", + "title": f"Issue {iid}", + "description": "x" * (tickets.MAX_TICKET_CHARACTERS + 1), + "labels": ["bug", "backend"], + } if outcome == 200 else {"message": "unavailable"} + response = requests.Response() + response.status_code = outcome + response._content = json.dumps(payload).encode() + response.headers["Content-Type"] = "application/json" + response.request = request + return response + + def close(self): + pass + + +@pytest.fixture +def provider_factory(monkeypatch): + # Keep unrelated integrations deterministic; retain the native extraction path. + monkeypatch.setattr(tickets, "add_jira_tickets", lambda provider, content: content) + sessions = [] + + def make(description, outcomes=(), token="offline-token", retry=False): + transport = IssueTransport(outcomes) + session = requests.Session() + session.mount("https://", transport) + sessions.append(session) + client = gitlab.Gitlab("https://gitlab.example/gitlab", private_token=token, session=session, + retry_transient_errors=retry) + provider = SimpleNamespace( + gl=client, id_project="group/repo", gitlab_url="https://gitlab.example/gitlab", + get_user_description=lambda: description, + supports_issue_url_tickets=lambda: False, + supports_issue_reference_tickets=lambda: True, + ) + return provider, transport + + yield make + for session in sessions: + session.close() + + +async def test_extract_issues_without_project_metadata(provider_factory): + provider, transport = provider_factory("Fixes #1 #2 other/sub/project#3 #4 #1") + result = await tickets.extract_tickets(provider) + + assert [ticket["ticket_id"] for ticket in result] == [1, 2, 3] + assert [urlparse(request.url).path for request in transport.requests] == [ + "/gitlab/api/v4/projects/group%2Frepo/issues/1", + "/gitlab/api/v4/projects/group%2Frepo/issues/2", + "/gitlab/api/v4/projects/other%2Fsub%2Fproject/issues/3", + ] + assert all(request.headers["PRIVATE-TOKEN"] == "offline-token" for request in transport.requests) + assert all(ticket["body"] == "x" * tickets.MAX_TICKET_CHARACTERS + "..." for ticket in result) + assert all(ticket["labels"] == "bug, backend" for ticket in result) + + +@pytest.mark.parametrize("failure", [401, 403, 404, 500, requests.Timeout("offline timeout")]) +async def test_issue_failures_preserve_successful_context(provider_factory, failure): + provider, transport = provider_factory("Fixes #1 #2", [failure, 200]) + result = await tickets.extract_tickets(provider) + + assert [ticket["ticket_id"] for ticket in result] == [2] + assert len(transport.requests) == 2 + + +async def test_failed_issue_is_fetched_again_on_next_extraction(provider_factory): + provider, transport = provider_factory("Fixes #1", [403, 200]) + assert await tickets.extract_tickets(provider) == [] + assert [ticket["ticket_id"] for ticket in await tickets.extract_tickets(provider)] == [1] + assert len(transport.requests) == 2 + + +async def test_gitlab_transport_retry_remains_enabled(provider_factory, monkeypatch): + monkeypatch.setattr("gitlab.utils.time.sleep", lambda seconds: None) + provider, transport = provider_factory("Fixes #1", [503, 200], retry=True) + assert [ticket["ticket_id"] for ticket in await tickets.extract_tickets(provider)] == [1] + assert len(transport.requests) == 2 + + +async def test_cancellation_propagates_without_fetching_later_issues(provider_factory): + provider, transport = provider_factory("Fixes #1 #2", [asyncio.CancelledError()]) + with pytest.raises(asyncio.CancelledError): + await tickets.extract_tickets(provider) + assert len(transport.requests) == 1 + + +async def test_separate_credentials_do_not_share_results(provider_factory): + allowed, allowed_transport = provider_factory("Fixes #1", token="allowed-offline-token") + denied, denied_transport = provider_factory("Fixes #1", [403], token="denied-offline-token") + assert len(await tickets.extract_tickets(allowed)) == 1 + assert await tickets.extract_tickets(denied) == [] + assert allowed_transport.requests[0].headers["PRIVATE-TOKEN"] == "allowed-offline-token" + assert denied_transport.requests[0].headers["PRIVATE-TOKEN"] == "denied-offline-token" diff --git a/tests/unittest/test_gunicorn_config.py b/tests/unittest/test_gunicorn_config.py index 153b658a6a..acb32efd08 100644 --- a/tests/unittest/test_gunicorn_config.py +++ b/tests/unittest/test_gunicorn_config.py @@ -1,3 +1,7 @@ +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + import pytest from pr_agent.servers import gunicorn_config @@ -174,3 +178,38 @@ def test_when_ready_freezes_gc(monkeypatch): monkeypatch.setattr(gunicorn_config.gc, "freeze", lambda: calls.append(True)) gunicorn_config.when_ready(server=None) assert calls == [True] + + +def _fake_prometheus_client(monkeypatch): + """Stub prometheus_client so the gate tests never import the real module. + + Importing the real module would let it read PROMETHEUS_MULTIPROC_DIR and + permanently decide the process-wide storage mode, which collides with the + telemetry suite's import-time decision. A fake module keeps the test + isolated and still exercises child_exit's own import plus the gate. + """ + multiproc_module = ModuleType("prometheus_client.multiprocess") + multiproc_module.mark_process_dead = Mock() + client_module = ModuleType("prometheus_client") + client_module.multiprocess = multiproc_module + monkeypatch.setitem(sys.modules, "prometheus_client", client_module) + monkeypatch.setitem(sys.modules, "prometheus_client.multiprocess", multiproc_module) + return multiproc_module + + +def test_child_exit_skips_prometheus_without_multiproc_dir(monkeypatch): + fake_multiprocess = _fake_prometheus_client(monkeypatch) + monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR", raising=False) + + gunicorn_config.child_exit(server=None, worker=SimpleNamespace(pid=123)) + + fake_multiprocess.mark_process_dead.assert_not_called() + + +def test_child_exit_marks_worker_dead_when_multiproc_dir_set(monkeypatch, tmp_path): + fake_multiprocess = _fake_prometheus_client(monkeypatch) + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + + gunicorn_config.child_exit(server=None, worker=SimpleNamespace(pid=456)) + + fake_multiprocess.mark_process_dead.assert_called_once_with(456) diff --git a/tests/unittest/test_jira_ticket_extraction.py b/tests/unittest/test_jira_ticket_extraction.py new file mode 100644 index 0000000000..48f54a5731 --- /dev/null +++ b/tests/unittest/test_jira_ticket_extraction.py @@ -0,0 +1,693 @@ +from string import ascii_uppercase +from unittest.mock import MagicMock, patch +from urllib.parse import urlparse + +import pytest + +from pr_agent.config_loader import get_settings +from pr_agent.git_providers import AzureDevopsProvider +from pr_agent.tools.ticket_pr_compliance_check import ( + MAX_JIRA_FETCH_ATTEMPTS, + MAX_TICKET_CHARACTERS, + MAX_TICKETS, + _get_jira_client, + _get_pr_title, + _jira_cloud_base_url, + add_jira_tickets, + extract_jira_tickets, + extract_tickets, + find_jira_tickets, +) + +# Keys the tests mutate via get_settings().set(...). Snapshot and restore them around +# every test so values (e.g. JIRA_REQUIREMENTS_FIELD) don't leak between tests. +_JIRA_KEYS = ( + "JIRA.JIRA_SITE", + "JIRA.JIRA_API_EMAIL", + "JIRA.JIRA_API_TOKEN", + "JIRA.JIRA_REQUIREMENTS_FIELD", + "JIRA.PROJECT_KEYS", +) + + +@pytest.fixture(autouse=True) +def restore_jira_settings(): + saved = {key: get_settings().get(key, None) for key in _JIRA_KEYS} + yield + for key, value in saved.items(): + get_settings().set(key, value) + + +class TestFindJiraTickets: + """Jira key extraction from arbitrary text (PR title, description, branch name).""" + + def test_uppercase_key_with_prefix(self): + """feature/ABC-123-description -> ABC-123""" + assert find_jira_tickets("feature/ABC-123-description-of-branch") == ["ABC-123"] + + def test_lowercase_key_with_prefix_normalized(self): + """bugfix/abc-123-description -> ABC-123 (case-insensitive, normalized to upper)""" + assert find_jira_tickets("bugfix/abc-123-description-of-branch") == ["ABC-123"] + + def test_mixed_case_key_normalized(self): + """Abc-123 -> ABC-123""" + assert find_jira_tickets("Abc-123-fix") == ["ABC-123"] + + def test_arbitrary_prefix_segment(self): + """Any prefix segment, not just feature/bugfix.""" + assert find_jira_tickets("chore/PROJ-45-cleanup") == ["PROJ-45"] + assert find_jira_tickets("hotfix/proj-45-cleanup") == ["PROJ-45"] + + def test_key_at_start_no_prefix(self): + """ABC-123-fix -> ABC-123""" + assert find_jira_tickets("ABC-123-fix") == ["ABC-123"] + + def test_key_anywhere_in_branch(self): + """Key embedded in the middle of a branch name is still found.""" + assert find_jira_tickets("release/v1.2.3-ABC-9-final") == ["ABC-9"] + + def test_key_in_description_text(self): + """Key mentioned in free text.""" + assert find_jira_tickets("This implements ABC-123 as discussed") == ["ABC-123"] + + def test_browse_url(self): + """Full Jira browse URL -> key.""" + assert find_jira_tickets( + "see https://acme.atlassian.net/browse/ABC-123 for details" + ) == ["ABC-123"] + + def test_no_ticket(self): + """Branch with no key -> []""" + assert find_jira_tickets("feature/no-ticket-here") == [] + assert find_jira_tickets("") == [] + + def test_multiple_tickets_deduped_in_order(self): + """Distinct keys are returned de-duplicated and case-normalized, in first-seen + order. Order must be stable so the later MAX_TICKETS cap is deterministic.""" + result = find_jira_tickets("ABC-1 and DEF-2, again ABC-1 and abc-1") + assert result == ["ABC-1", "DEF-2"] + + def test_order_preserved_across_patterns(self): + """First-seen order holds even when keys arrive via different patterns (plain + key vs. browse URL).""" + result = find_jira_tickets( + "GHI-3 first, then https://acme.atlassian.net/browse/ABC-1, then DEF-2" + ) + assert result == ["GHI-3", "ABC-1", "DEF-2"] + + +class TestExtractJiraTickets: + """End-to-end extraction: find keys, fetch via the Jira client, map to ticket dicts.""" + + def _configure_jira(self): + get_settings().set("JIRA.JIRA_SITE", "acme") + get_settings().set("JIRA.JIRA_API_EMAIL", "me@acme.com") + get_settings().set("JIRA.JIRA_API_TOKEN", "token123") + + def _disable_jira(self): + get_settings().set("JIRA.JIRA_SITE", "") + get_settings().set("JIRA.JIRA_API_EMAIL", "") + get_settings().set("JIRA.JIRA_API_TOKEN", "") + + def _fake_client(self, fields=None): + fields = fields or {"summary": "Title", "description": "Body", "labels": ["a"]} + client = MagicMock() + client.issue.return_value = {"fields": fields} + return client + + def test_returns_empty_when_not_configured(self): + self._disable_jira() + assert extract_jira_tickets("bugfix/abc-123-x") == [] + + def test_no_client_built_when_no_keys(self): + """No Jira keys in the text -> return early without constructing a client, so a + keyless PR pays no client-init cost (or noisy init-failure log).""" + self._configure_jira() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls: + result = extract_jira_tickets("nothing ticket-like here") + assert result == [] + jira_cls.assert_not_called() + + def test_fetches_lowercase_branch_key(self): + """The whole point: a lowercased branch key is detected and fetched as upper.""" + self._configure_jira() + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("bugfix/abc-123-description-of-branch") + client.issue.assert_called_once_with("ABC-123") + assert len(result) == 1 + assert result[0]["ticket_id"] == "ABC-123" + assert result[0]["ticket_url"] == "https://acme.atlassian.net/browse/ABC-123" + assert result[0]["title"] == "Title" + assert result[0]["labels"] == "a" + + def test_requirements_field_populated_when_configured(self): + """When jira_requirements_field is set, that custom field maps to requirements.""" + self._configure_jira() + get_settings().set("JIRA.JIRA_REQUIREMENTS_FIELD", "customfield_10127") + client = MagicMock() + client.issue.return_value = {"fields": { + "summary": "T", "description": "B", "labels": [], + "customfield_10127": "Acceptance criteria text", + }} + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("ABC-1") + assert result[0]["requirements"] == "Acceptance criteria text" + + def test_requirements_truncated_like_body(self): + """A large requirements custom field is capped the same way the body is, so it + can't push an unbounded blob into the review prompt.""" + self._configure_jira() + get_settings().set("JIRA.JIRA_REQUIREMENTS_FIELD", "customfield_10127") + client = MagicMock() + client.issue.return_value = {"fields": { + "summary": "T", "description": "B", "labels": [], + "customfield_10127": "x" * 50, + }} + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("ABC-1", max_characters=10) + assert result[0]["requirements"] == "x" * 10 + "..." + + def test_requirements_empty_when_field_not_configured(self): + """With no requirements field configured, requirements stays empty.""" + self._configure_jira() + get_settings().set("JIRA.JIRA_REQUIREMENTS_FIELD", "") + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("ABC-1") + assert result[0]["requirements"] == "" + + def test_checks_all_tickets_when_multiple(self): + """When several distinct keys are present, each is fetched.""" + self._configure_jira() + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("ABC-1 DEF-2 GHI-3") + fetched = {c.args[0] for c in client.issue.call_args_list} + assert fetched == {"ABC-1", "DEF-2", "GHI-3"} + assert len(result) == 3 + + def test_caps_resolved_tickets_not_candidates(self): + """The cap counts tickets that resolve, so it stops the fetch loop rather than + truncating the candidate list.""" + self._configure_jira() + client = self._fake_client() + keys = " ".join(f"ABC-{i}" for i in range(1, MAX_TICKETS + 3)) + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets(keys) + assert client.issue.call_count == MAX_TICKETS + assert len(result) == MAX_TICKETS + + def test_noise_keys_do_not_displace_a_real_ticket(self): + """SHA-1, SHA-256, UTF-8 and AES-256 all match the key pattern. Capping candidates + before the fetch spends the budget on them and drops the ticket the PR names.""" + self._configure_jira() + + def fake_issue(key): + if key == "PROJ-4242": + return {"fields": {"summary": "Real", "description": "Body", "labels": []}} + raise Exception("404 not found") + + client = MagicMock() + client.issue.side_effect = fake_issue + text = ("Switch the digest from SHA-1 to SHA-256, normalise to UTF-8 and rotate " + "the AES-256 key. Implements PROJ-4242") + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets(text) + assert [t["ticket_id"] for t in result] == ["PROJ-4242"] + + def test_real_ticket_reached_past_more_noise_than_the_cap(self): + """The same property with more noise keys than MAX_TICKETS, so it holds whatever + the cap is set to.""" + self._configure_jira() + + def fake_issue(key): + if key == "PROJ-4242": + return {"fields": {"summary": "Real", "description": "Body", "labels": []}} + raise Exception("404 not found") + + client = MagicMock() + client.issue.side_effect = fake_issue + noise = " ".join(f"AES-{i}" for i in range(1, MAX_TICKETS + 3)) + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets(f"{noise} PROJ-4242") + assert [t["ticket_id"] for t in result] == ["PROJ-4242"] + + def test_fetch_attempts_are_bounded(self): + """Key-shaped noise cannot make one PR pay unbounded authenticated Jira calls.""" + self._configure_jira() + client = MagicMock() + client.issue.side_effect = Exception("404 not found") + keys = " ".join(f"ABC-{i}" for i in range(1, MAX_JIRA_FETCH_ATTEMPTS + 6)) + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets(keys) + assert client.issue.call_count == MAX_JIRA_FETCH_ATTEMPTS + assert result == [] + + def test_skips_ticket_on_fetch_error(self): + """A failed fetch for one key does not abort the others.""" + self._configure_jira() + client = MagicMock() + client.issue.side_effect = [ + Exception("404 not found"), + {"fields": {"summary": "Second", "description": "B", "labels": []}}, + ] + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("ABC-1 ABC-2") + assert len(result) == 1 + assert result[0]["title"] == "Second" + + def test_nonexistent_keys_are_skipped(self): + """Key-like noise (utf-8, sha-1) that does not resolve in Jira is skipped, + leaving only the real ticket.""" + self._configure_jira() + + def fake_issue(key): + if key == "ABC-123": + return {"fields": {"summary": "Real", "description": "Body", "labels": []}} + raise Exception("404 not found") + + client = MagicMock() + client.issue.side_effect = fake_issue + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("abc-123 utf-8") + assert [t["ticket_id"] for t in result] == ["ABC-123"] + + def test_project_keys_allowlist_filters_noise_before_fetch(self): + """When jira.project_keys is set, key-shaped noise outside the allowlist is + dropped before any lookup: it never reaches jira_client.issue().""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", ["PROJ"]) + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("Switch SHA-256 digest per PROJ-4242") + client.issue.assert_called_once_with("PROJ-4242") + assert [t["ticket_id"] for t in result] == ["PROJ-4242"] + + def test_project_keys_allowlist_drops_unlisted_prefixes_before_lookup(self): + """With several projects listed, every listed prefix is kept in first-seen order + and every other prefix ("SHA-256", "UTF-8", "ISO-8601") costs no authenticated 404.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", ["PROJ", "OPS"]) + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("PROJ-1 uses SHA-256 and UTF-8, see OPS-7 and ISO-8601") + assert [call.args[0] for call in client.issue.call_args_list] == ["PROJ-1", "OPS-7"] + assert [t["ticket_id"] for t in result] == ["PROJ-1", "OPS-7"] + + def test_project_keys_allowlist_empty_result_skips_client(self): + """When every candidate is filtered out by the allowlist, no Jira client is + built at all: same early-exit as having no keys in the text.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", ["PROJ"]) + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls: + result = extract_jira_tickets("SHA-256 and UTF-8 only, no real ticket") + assert result == [] + jira_cls.assert_not_called() + + def test_uppercase_allowlist_still_matches_a_lowercased_ticket_key(self): + """Detected keys are normalized to upper case before the allowlist is applied, so a + lowercased branch name (bugfix/proj-12-x) matches an upper-case entry.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", ["PROJ"]) + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("bugfix/proj-12-x mentions SHA-256") + assert [call.args[0] for call in client.issue.call_args_list] == ["PROJ-12"] + assert [t["ticket_id"] for t in result] == ["PROJ-12"] + + @pytest.mark.parametrize("entries", [["proj"], ["Proj"], "proj, ops"]) + def test_lowercase_project_keys_fail_closed(self, entries): + """Jira project keys are upper case, so a lower-case entry is a typo rather than a + spelling variant. Accepting it would allow a project the repo never named, so it is + rejected with a warning and no lookup happens at all.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", entries) + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls, \ + patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_logger: + result = extract_jira_tickets("bugfix/proj-12-x mentions SHA-256") + assert result == [] + jira_cls.assert_not_called() + warning_calls = [str(c) for c in get_logger.return_value.warning.call_args_list] + assert any("project_keys" in c and "index 0" in c for c in warning_calls) + assert any("no valid entry" in c for c in warning_calls) + + def test_project_keys_accept_comma_separated_string(self): + """Environment-variable overrides arrive as a string (jira__project_keys="PROJ, OPS").""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", "PROJ, OPS") + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + extract_jira_tickets("OPS-3 SHA-256 PROJ-4") + assert [call.args[0] for call in client.issue.call_args_list] == ["OPS-3", "PROJ-4"] + + @pytest.mark.parametrize("bad_entry", [ + "https://acme.atlassian.net/browse/PROJ", # a URL, not a key + "PROJ-123", # a ticket key, not a project key + "P", # too short to ever match a ticket + "PROJECT_KEY_TOO_LONG", # longer than the ticket pattern allows + "pro j", # whitespace + ]) + def test_invalid_project_keys_are_ignored_with_a_warning(self, bad_entry): + """A malformed entry cannot widen the allowlist; it is dropped with a warning while + the valid entries keep filtering.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", [bad_entry, "PROJ"]) + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client), \ + patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_logger: + extract_jira_tickets("PROJ-1 SHA-256") + assert [call.args[0] for call in client.issue.call_args_list] == ["PROJ-1"] + warning_calls = [str(c) for c in get_logger.return_value.warning.call_args_list] + # Named by position, not quoted back: see + # test_a_rejected_entry_is_never_echoed_into_the_log. + assert any("project_keys" in c and "index 0" in c for c in warning_calls) + + def test_a_rejected_entry_is_never_echoed_into_the_log(self): + """The warning names the entry's position and type, never its content. + + project_keys is operator configuration this tool does not control. A value pasted + into the wrong setting can be a credential, and one carrying newlines could forge + whole log lines, so the rejected text is not quoted back. + """ + self._configure_jira() + secret = "hunter2-SECRET\nWARNING: Jira ticket lookup disabled by operator" + get_settings().set("JIRA.PROJECT_KEYS", [secret, "PROJ"]) + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client), \ + patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_logger: + extract_jira_tickets("PROJ-1 SHA-256") + # The valid entry still filters, so the rejection is not a silent drop. + assert [call.args[0] for call in client.issue.call_args_list] == ["PROJ-1"] + warnings = [str(c) for c in get_logger.return_value.warning.call_args_list] + assert any("project_keys" in w and "index 0" in w and "type str" in w for w in warnings) + for w in warnings: + assert "hunter2" not in w + assert "SECRET" not in w + assert "\n" not in w + assert "disabled by operator" not in w + + def test_filtering_keeps_order_and_reports_every_skipped_key(self): + """Candidates are partitioned in a single pass over the list. + + The candidate list is not capped until lookups begin, so a PR carrying many + distinct key-shaped tokens is the shape this has to stay linear on. What the + caller sees must not change: kept keys keep their first-seen order and the debug + line still names every skipped key. + """ + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", ["PROJ"]) + noise = [f"X{letter}-{index}" for index, letter in enumerate(ascii_uppercase, start=1)] + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client), \ + patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_logger: + extract_jira_tickets(" ".join(["PROJ-1", *noise, "PROJ-2"])) + assert [call.args[0] for call in client.issue.call_args_list] == ["PROJ-1", "PROJ-2"] + debug_calls = [str(c) for c in get_logger.return_value.debug.call_args_list] + assert any( + "Skipping Jira lookup for keys outside jira.project_keys" in c + and ", ".join(noise) in c + for c in debug_calls + ) + + @pytest.mark.parametrize("entries", [["PROJ-123"], ["https://acme.atlassian.net/browse/PROJ", "P"], "PROJ-123, x"]) + def test_project_keys_with_no_valid_entry_fail_closed(self, entries): + """A configured allowlist whose entries are all malformed must not widen to the + unfiltered default: nothing is looked up and no client is built until it is fixed.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", entries) + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls, \ + patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_logger: + result = extract_jira_tickets("PROJ-1 SHA-256") + assert result == [] + jira_cls.assert_not_called() + warning_calls = [str(c) for c in get_logger.return_value.warning.call_args_list] + assert any("no valid entry" in c for c in warning_calls) + + @pytest.mark.parametrize("value", [True, False, 0, 1, 1.5, {"PROJ": True}]) + def test_project_keys_of_an_unsupported_type_fail_closed(self, value): + """`project_keys=true` / `=false` from a YAML or CLI override is neither unset nor a + list; it must not raise and must not re-enable lookups for every key.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", value) + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls, \ + patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_logger: + result = extract_jira_tickets("PROJ-1 SHA-256") + assert result == [] + jira_cls.assert_not_called() + warning_calls = [str(c) for c in get_logger.return_value.warning.call_args_list] + assert any("must be a list" in c for c in warning_calls) + + @pytest.mark.parametrize("bad_entry", [True, False, None, 0]) + def test_non_string_project_keys_are_not_stringified_into_labels(self, bad_entry): + """A boolean or null in the list must not become a key-shaped label ("TRUE", + "NONE") that silently filters; it is ignored and the valid entry keeps working.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", [bad_entry, "PROJ"]) + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client), \ + patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_logger: + extract_jira_tickets("PROJ-1 TRUE-2 NONE-3 SHA-256") + assert [call.args[0] for call in client.issue.call_args_list] == ["PROJ-1"] + warning_calls = [str(c) for c in get_logger.return_value.warning.call_args_list] + assert any("project_keys" in c for c in warning_calls) + + @pytest.mark.parametrize("value", [[], (), "", " "]) + def test_unsupplied_project_keys_keep_looking_up_every_key(self, value): + """The three ways of supplying nothing: the option missing (covered by the tests + above), the shipped empty list, and a blank top-level string, since + `jira__project_keys=""` is how a shell spells an unset environment variable. All + keep today's behaviour.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", value) + client = self._fake_client() + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + result = extract_jira_tickets("PROJ-1 SHA-256") + assert [call.args[0] for call in client.issue.call_args_list] == ["PROJ-1", "SHA-256"] + assert len(result) == 2 + + @pytest.mark.parametrize("entries", [[""], [" "], ["", ""], ",", " , "]) + def test_supplied_blank_project_keys_fail_closed(self, entries): + """A blank *entry* is a malformed override, not a second spelling of the default: it + carries a list or a separator, so something was supplied and got lost. Widening back + to unfiltered lookup there would let a broken override grant access it never named.""" + self._configure_jira() + get_settings().set("JIRA.PROJECT_KEYS", entries) + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls, \ + patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_logger: + result = extract_jira_tickets("PROJ-1 SHA-256") + assert result == [] + jira_cls.assert_not_called() + warning_calls = [str(c) for c in get_logger.return_value.warning.call_args_list] + assert any("no valid entry" in c for c in warning_calls) + + +class TestGetJiraClient: + """Cloud client construction: site-name -> base URL, email+token auth, v2 pin.""" + + def _set(self, site=None, email=None, token=None): + get_settings().set("JIRA.JIRA_SITE", site if site is not None else "") + get_settings().set("JIRA.JIRA_API_EMAIL", email if email is not None else "") + get_settings().set("JIRA.JIRA_API_TOKEN", token if token is not None else "") + + def test_cloud_builds_url_from_site_and_pins_v2(self): + """site name -> https://.atlassian.net, email/token basic auth, v2 pinned.""" + self._set(site="acme", email="me@acme.com", token="token123") + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls: + _get_jira_client() + jira_cls.assert_called_once_with( + url="https://acme.atlassian.net", username="me@acme.com", + password="token123", api_version="2", + ) + + def test_returns_none_when_not_configured(self): + self._set() + assert _get_jira_client() is None + + def test_returns_none_when_email_missing(self): + """Cloud requires email; site + token alone is incomplete.""" + self._set(site="acme", token="token123") + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls: + assert _get_jira_client() is None + jira_cls.assert_not_called() + + def test_no_warning_when_nothing_configured(self): + """Jira simply not in use -> return None silently, no misconfiguration warning.""" + self._set() + with patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_log: + assert _get_jira_client() is None + get_log.return_value.warning.assert_not_called() + + def test_warns_when_partially_configured(self): + """Some [jira] value set but the required site + email + token are incomplete + -> warn (likely a misconfiguration) and return None.""" + self._set(email="me@acme.com", token="token123") # site missing + with patch("pr_agent.tools.ticket_pr_compliance_check.get_logger") as get_log: + assert _get_jira_client() is None + get_log.return_value.warning.assert_called_once() + msg = get_log.return_value.warning.call_args.args[0] + assert "jira_site" in msg + assert "jira_api_token" not in msg # the one that IS set is not listed as missing + + +class TestJiraSiteInjection: + """jira_site is validated so repo-controlled config can't redirect the authenticated + request (and token) off *.atlassian.net. Building the URL from the site name means a + malicious value cannot express a different host.""" + + # Values a malicious repo config might inject to try to escape *.atlassian.net. + INJECTION_ATTEMPTS = [ + "evil.website.com?", # query separator + "evil.com#", # fragment + "evil.com/", # path separator + "evil.com", # dotted host + "evil.com:443", # port + "x@evil.com", # userinfo + "a.evil.com", # subdomain prefix + "../../evil", # path traversal + "evil%2ecom", # encoded dot + "evil_underscore", # underscore not allowed in DNS label + "evil .com", # internal whitespace + "", # empty + ] + + @pytest.mark.parametrize("bad_site", INJECTION_ATTEMPTS) + def test_injection_attempt_rejected(self, bad_site): + get_settings().set("JIRA.JIRA_SITE", bad_site) + get_settings().set("JIRA.JIRA_API_EMAIL", "me@acme.com") + get_settings().set("JIRA.JIRA_API_TOKEN", "token123") + # No client is built, and no base URL is produced for an invalid site. + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls: + assert _get_jira_client() is None + jira_cls.assert_not_called() + assert _jira_cloud_base_url() is None + + def test_valid_site_only_ever_targets_atlassian_net(self): + """Any accepted site name resolves to a host under .atlassian.net.""" + for good in ("acme", "my-org", "a1b2", "x"): + get_settings().set("JIRA.JIRA_SITE", good) + url = _jira_cloud_base_url() + assert url == f"https://{good}.atlassian.net" + assert urlparse(url).hostname.endswith(".atlassian.net") + + +class TestGetPrTitle: + """Provider-agnostic title access (GitHub/Bitbucket use .pr, GitLab uses .mr).""" + + def test_reads_pr_title(self): + gp = MagicMock(spec=["pr"]) + gp.pr = MagicMock(title="From PR object") + assert _get_pr_title(gp) == "From PR object" + + def test_reads_mr_title_when_no_pr(self): + """GitLab stores the merge request as .mr, not .pr.""" + gp = MagicMock(spec=["mr"]) + gp.mr = MagicMock(title="From MR object") + assert _get_pr_title(gp) == "From MR object" + + def test_returns_empty_when_no_title(self): + gp = MagicMock(spec=[]) + assert _get_pr_title(gp) == "" + + +class TestAddJiraTickets: + """Provider-agnostic Jira append used by extract_tickets for every provider.""" + + def _provider(self, title="", description="", branch=""): + gp = MagicMock(spec=["pr", "get_user_description", "get_pr_branch"]) + gp.pr = MagicMock(title=title) + gp.get_user_description.return_value = description + gp.get_pr_branch.return_value = branch + return gp + + def _configure_jira(self): + get_settings().set("JIRA.JIRA_SITE", "acme") + get_settings().set("JIRA.JIRA_API_EMAIL", "me@acme.com") + get_settings().set("JIRA.JIRA_API_TOKEN", "token123") + get_settings().set("JIRA.JIRA_REQUIREMENTS_FIELD", "") + + def test_appends_ticket_from_any_provider(self): + """Works off get_user_description + get_pr_branch, so it is provider-neutral.""" + self._configure_jira() + client = MagicMock() + client.issue.return_value = {"fields": {"summary": "T", "description": "B", "labels": []}} + gp = self._provider(branch="feature/ABC-123-x") + out = [] + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + add_jira_tickets(gp, out) + assert [t["ticket_id"] for t in out] == ["ABC-123"] + + def test_dedupes_against_existing_tickets(self): + """A Jira ticket already present (same url) is not added twice.""" + self._configure_jira() + client = MagicMock() + client.issue.return_value = {"fields": {"summary": "T", "description": "B", "labels": []}} + gp = self._provider(title="ABC-123") + existing = [{"ticket_url": "https://acme.atlassian.net/browse/ABC-123"}] + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + add_jira_tickets(gp, existing) + assert len(existing) == 1 + + def test_noop_when_jira_not_configured(self): + get_settings().set("JIRA.JIRA_SITE", "") + get_settings().set("JIRA.JIRA_API_TOKEN", "") + gp = self._provider(branch="feature/ABC-123-x") + out = [] + add_jira_tickets(gp, out) + assert out == [] + + def test_respects_overall_cap_with_existing_tickets(self): + """MAX_TICKETS is the combined per-PR cap: provider-native tickets already in + tickets_content count against it, and Jira is appended only up to the cap.""" + self._configure_jira() + client = MagicMock() + client.issue.return_value = {"fields": {"summary": "T", "description": "B", "labels": []}} + # Pre-fill with (MAX_TICKETS - 1) provider-native tickets, then offer several Jira keys. + existing = [{"ticket_url": f"https://example/issues/{i}"} for i in range(MAX_TICKETS - 1)] + gp = self._provider(description="ABC-1 DEF-2 GHI-3 JKL-4") + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira", return_value=client): + add_jira_tickets(gp, existing) + assert len(existing) == MAX_TICKETS # only one Jira ticket was appended + + def test_skips_jira_entirely_when_cap_already_reached(self): + """When provider-native tickets already fill the cap, no Jira client is built.""" + self._configure_jira() + existing = [{"ticket_url": f"https://example/issues/{i}"} for i in range(MAX_TICKETS)] + gp = self._provider(description="ABC-1 DEF-2") + with patch("pr_agent.tools.ticket_pr_compliance_check.Jira") as jira_cls: + add_jira_tickets(gp, existing) + jira_cls.assert_not_called() + assert len(existing) == MAX_TICKETS + + +class TestAzureRequirementsTruncation: + """The Azure DevOps branch caps acceptance criteria like the body (no unbounded blob).""" + + @pytest.mark.asyncio + async def test_acceptance_criteria_truncated(self): + gp = AzureDevopsProvider.__new__(AzureDevopsProvider) # bypass __init__/network + gp.get_linked_work_items = MagicMock(return_value=[{ + "id": 1, + "url": "https://dev.azure.com/org/proj/_workitems/edit/1", + "title": "Work item", + "body": "short body", + "acceptance_criteria": "x" * (MAX_TICKET_CHARACTERS + 50), + "labels": [], + }]) + # Isolate the Azure branch: keep the provider-agnostic Jira step a no-op. + with patch("pr_agent.tools.ticket_pr_compliance_check.add_jira_tickets", + side_effect=lambda gp, tc: tc): + result = await extract_tickets(gp) + assert result[0]["requirements"] == "x" * MAX_TICKET_CHARACTERS + "..." + + @pytest.mark.asyncio + async def test_non_string_acceptance_criteria_becomes_empty(self): + gp = AzureDevopsProvider.__new__(AzureDevopsProvider) + gp.get_linked_work_items = MagicMock(return_value=[{ + "id": 2, "url": "u", "title": "t", "body": "b", + "acceptance_criteria": {"unexpected": "dict"}, "labels": [], + }]) + with patch("pr_agent.tools.ticket_pr_compliance_check.add_jira_tickets", + side_effect=lambda gp, tc: tc): + result = await extract_tickets(gp) + assert result[0]["requirements"] == "" diff --git a/tests/unittest/test_language_handler.py b/tests/unittest/test_language_handler.py index 400bd21d4f..fcc185dbe6 100644 --- a/tests/unittest/test_language_handler.py +++ b/tests/unittest/test_language_handler.py @@ -64,6 +64,28 @@ def test_edge_case_empty_languages(self): expected_output = [{'language': 'Other', 'files': files}] assert sort_files_by_main_languages(languages, files) == expected_output + # PyGithub 2.x adds a "url" string to the languages dict; it must not take part in the ranking + def test_non_numeric_language_entries_are_ignored(self): + languages = {'Python': 10, 'Java': 5, 'url': 'https://api.github.com/repos/o/r/languages'} + files = [ + type('', (object,), {'filename': 'file1.py'})(), + type('', (object,), {'filename': 'file2.java'})() + ] + expected_output = [ + {'language': 'Python', 'files': [files[0]]}, + {'language': 'Java', 'files': [files[1]]}, + {'language': 'Other', 'files': []} + ] + assert sort_files_by_main_languages(languages, files) == expected_output + + def test_main_pr_language_ignores_non_numeric_entries(self): + from pr_agent.git_providers.git_provider import get_main_pr_language + + languages = {'Python': 10, 'Java': 5, 'url': 'https://api.github.com/repos/o/r/languages'} + files = [type('', (object,), {'filename': 'file1.py'})()] + assert get_main_pr_language(languages, files) == 'python' + assert get_main_pr_language({'url': 'https://api.github.com/repos/o/r/languages'}, files) == '' + # Tests that function handles empty files list def test_edge_case_empty_files(self): languages = {'Python': 10, 'Java': 5} diff --git a/tests/unittest/test_linked_ticket_token_budget.py b/tests/unittest/test_linked_ticket_token_budget.py index d855c5f554..a54acb7647 100644 --- a/tests/unittest/test_linked_ticket_token_budget.py +++ b/tests/unittest/test_linked_ticket_token_budget.py @@ -204,6 +204,8 @@ def get_diff(_provider, token_handler, _model, **_kwargs): async def get_prediction(_model, *_args, **_kwargs): rendered_vars.append(tool.vars) + if tool_name == "review": + return "review:\n summary: prediction" return "prediction" monkeypatch.setattr(module, "fit_related_tickets_to_prompt_budget", fit_payload) @@ -258,7 +260,8 @@ def get_multiple_patches(_provider, token_handler, _model): async def get_prediction(_model, patches_diff=None, prompt=None): prediction_calls.append((prompt, patches_diff, tool.vars)) if prompt == "pr_description_only_files_prompts": - return "pr_files:\n- filename: src/app.py" + return ("pr_files:\n- filename: src/app.py\n changes_title: Update app\n" + " changes_summary: Updates app behavior.\n label: enhancement") return "title: Test\ndescription: Test" monkeypatch.setattr(module, "fit_related_tickets_to_prompt_budget", fit_payload) diff --git a/tests/unittest/test_litellm_api_key_guard.py b/tests/unittest/test_litellm_api_key_guard.py index a477d45960..f05ec35703 100644 --- a/tests/unittest/test_litellm_api_key_guard.py +++ b/tests/unittest/test_litellm_api_key_guard.py @@ -1,24 +1,27 @@ -""" -Tests for the litellm.api_key guard in LiteLLMAIHandler.chat_completion. - -Verifies: - - Placeholder key (DUMMY_LITELLM_API_KEY) is never injected into the call, and never - occupies the global litellm.api_key (issue #2544). - - None is not injected (e.g. when OpenAI key is set via litellm.openai_key). - - Real provider keys (Groq, SambaNova, XAI, OpenRouter, Azure AD) ARE injected. -""" -import contextlib +"""Tests for request-local LiteLLM provider credentials and endpoints.""" + +import asyncio +import atexit +import inspect +import json +import os +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch +from urllib.parse import unquote, urlparse +import httpx import litellm +import openai import pytest import pr_agent.algo.ai_handlers.litellm_ai_handler as litellm_handler +import pr_agent.algo.ai_handlers.litellm_helpers as litellm_helpers from pr_agent.algo.ai_handlers.litellm_ai_handler import DUMMY_LITELLM_API_KEY, LiteLLMAIHandler -def _make_settings(): - """Minimal settings object that satisfies __init__ and chat_completion.""" +def _make_settings(overrides=None): + overrides = overrides or {} return type("Settings", (), { "config": type("Config", (), { "reasoning_effort": None, @@ -30,14 +33,14 @@ def _make_settings(): "get": lambda self, key, default=None: default, })(), "litellm": type("LiteLLM", (), { + "extra_headers": overrides.get("LITELLM.EXTRA_HEADERS"), "get": lambda self, key, default=None: default, })(), - "get": lambda self, key, default=None: default, + "get": lambda self, key, default=None: overrides.get(key, default), })() def _mock_response(): - """Minimal acompletion response.""" mock = MagicMock() mock.__getitem__ = lambda self, key: { "choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}] @@ -47,650 +50,8157 @@ def _mock_response(): @pytest.fixture(autouse=True) -def patch_settings(monkeypatch): +def isolate_provider_state(monkeypatch): + from google.auth import pluggable + from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig + from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig + from litellm.llms.ragflow.chat.transformation import RAGFlowConfig + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + from litellm.llms.xai.chat.transformation import XAIChatConfig + + for config in (OpenAIGPTConfig, AzureAIStudioConfig, RAGFlowConfig, XAIChatConfig): + monkeypatch.setattr(config, "validate_environment", config.validate_environment) + executable_descriptors = { + name: pluggable.Credentials.__dict__.get(name) + for name in ("from_info", "refresh", "expired", "token_state") + } + monkeypatch.setattr(VertexBase, "load_auth", VertexBase.load_auth) + monkeypatch.setattr(VertexBase, "_credentials_from_service_account", VertexBase._credentials_from_service_account) + anthropic_get_api_key = inspect.getattr_static(litellm_handler.AnthropicModelInfo, "get_api_key") + anthropic_get_auth_token = inspect.getattr_static(litellm_handler.AnthropicModelInfo, "get_auth_token") + module_get_api_key = litellm_handler._anthropic_get_api_key + module_get_auth_token = litellm_handler._anthropic_get_auth_token + bedrock_mantle_auth_mixin = litellm_handler.BedrockMantleAuthMixin + module_bedrock_mantle_sign_request = litellm_handler._bedrock_mantle_sign_request + module_bedrock_mantle_resolve_bearer_token = litellm_handler._bedrock_mantle_resolve_bearer_token + if bedrock_mantle_auth_mixin is not None: + bedrock_mantle_sign_request = inspect.getattr_static( + bedrock_mantle_auth_mixin, + "sign_request", + ) + bedrock_mantle_resolve_bearer_token = inspect.getattr_static( + bedrock_mantle_auth_mixin, + "_resolve_bearer_token", + ) + provider_env_vars = { + variable + for variables in litellm_handler.PROVIDER_API_KEY_ENV_VARS.values() + for variable in variables + } | { + config.api_key_env + for provider in litellm_handler.JSONProviderRegistry.list_providers() + if ( + (config := litellm_handler.JSONProviderRegistry.get(provider)) is not None + and config.api_key_env + ) + } | { + variable + for variables in litellm_handler.PROVIDER_API_BASE_ENV_VARS.values() + for variable in variables + } | { + variable + for provider_variables in litellm_handler.PROVIDER_ROUTING_ENV_VARS.values() + for variables in provider_variables.values() + for variable in variables + } | { + config.api_base_env + for provider in litellm_handler.JSONProviderRegistry.list_providers() + if ( + (config := litellm_handler.JSONProviderRegistry.get(provider)) is not None + and config.api_base_env + ) + } | set(litellm_handler.AWS_CREDENTIAL_CHAIN_ENV_VARS) | { + "AWS_USE_IMDS", + "ARK_API_KEY", + "SNOWFLAKE_ACCOUNT_ID", + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_DEFAULT_REGION", + "AWS_BEDROCK_RUNTIME_ENDPOINT", + "BEDROCK_MANTLE_REGION", + "ANTHROPIC_AUTH_TOKEN", + "PALM_API_KEY", + "MOONSHOT_API_BASE", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "AZURE_API_BASE", + "AZURE_OPENAI_ENDPOINT", + "AZURE_API_VERSION", + "AZURE_AD_TOKEN", + "AZURE_OPENAI_AD_TOKEN", + "OPENROUTER_API_BASE", + "VERTEXAI_PROJECT", + "VERTEXAI_LOCATION", + "VERTEXAI_CREDENTIALS", + "GOOGLE_APPLICATION_CREDENTIALS", + "CLOUDSDK_CONFIG", + "CLOUDSDK_CORE_PROJECT", + "CLOUDSDK_ACTIVE_CONFIG_NAME", + "GOOGLE_CLOUD_QUOTA_PROJECT", + "APPENGINE_RUNTIME", + "OPENAI_PROJECT_ID", + "OPENAI_ORG_ID", + "OPENAI_CUSTOM_HEADERS", + "EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", + } + for variable in provider_env_vars: + monkeypatch.delenv(variable, raising=False) + provider_global_names = { + name + for names in litellm_handler.PROVIDER_API_KEY_GLOBALS.values() + for name in names + } + for name in provider_global_names | { + "api_key", + "openai_key", + "api_base", + "api_version", + "headers", + "organization", + "vertex_project", + "vertex_location", + }: + monkeypatch.setattr(litellm, name, None, raising=False) + monkeypatch.setattr(openai, "api_key", None) monkeypatch.setattr(litellm_handler, "get_settings", _make_settings) + yield + for name, descriptor in executable_descriptors.items(): + if descriptor is not None: + setattr(pluggable.Credentials, name, descriptor) + elif name in pluggable.Credentials.__dict__: + delattr(pluggable.Credentials, name) + litellm_handler.AnthropicModelInfo.get_api_key = anthropic_get_api_key + litellm_handler.AnthropicModelInfo.get_auth_token = anthropic_get_auth_token + litellm_handler._anthropic_get_api_key = module_get_api_key + litellm_handler._anthropic_get_auth_token = module_get_auth_token + if bedrock_mantle_auth_mixin is not None: + bedrock_mantle_auth_mixin.sign_request = bedrock_mantle_sign_request + bedrock_mantle_auth_mixin._resolve_bearer_token = bedrock_mantle_resolve_bearer_token + litellm_handler._bedrock_mantle_sign_request = module_bedrock_mantle_sign_request + litellm_handler._bedrock_mantle_resolve_bearer_token = module_bedrock_mantle_resolve_bearer_token -@pytest.fixture(autouse=True) -def restore_litellm_globals(): - """Undo the LiteLLM globals LiteLLMAIHandler.__init__ writes. - - Constructing the handler mutates process-wide state (litellm.api_key and - litellm.openai_key). monkeypatch only reverts what a test set itself, so without - this the placeholder written during __init__ would outlive the test and make - later tests order-dependent. - """ - saved = {name: getattr(litellm, name) for name in ("api_key", "openai_key")} - yield - for name, value in saved.items(): - setattr(litellm, name, value) +async def _assert_native_azure_ad_auth( + monkeypatch, transport, entrypoint, source, token_variable, *, captured_key=None, auth_headers=None, + expected_authorization="Bearer request-ad-token", expected_api_key=None, sdk_alias=None, alias_mutation="replace", + request_count=1, refresh_tokens=False, concurrent_handlers=False, +): + from litellm.caching.caching import DualCache + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.litellm_core_utils import logging_worker + from litellm.llms.azure import common_utils as azure_common + endpoint = "https://request.openai.azure.com" + if transport == "cloudflare": + endpoint = "https://gateway.ai.cloudflare.com/v1/account/gateway/azure-openai/resource" + settings = _make_settings({ + "OPENAI.API_TYPE": "azure", + "OPENAI.KEY": captured_key, + "OPENAI.API_BASE": endpoint, + "OPENAI.API_VERSION": "v1" if transport == "v1" else "2024-06-01", + "AZURE_AD.CLIENT_ID": "request-client" if source == "client_credential" else None, + "LITELLM.EXTRA_HEADERS": json.dumps({"x-request-header": "owned", **(auth_headers or {})}), + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_credential", lambda settings: object()) + tokens = iter(f"request-ad-token-{index}" for index in range(request_count)) + monkeypatch.setattr( + litellm_handler, "_get_azure_ad_token", + lambda credential: next(tokens) if refresh_tokens else "request-ad-token", + ) + if token_variable: + monkeypatch.setenv(token_variable, "oidc/request" if source == "oidc" else "request-ad-token") + if sdk_alias is not None: + monkeypatch.setenv("AZURE_OPENAI_AD_TOKEN", sdk_alias) + monkeypatch.setenv("AZURE_CLIENT_ID", "request-client") + monkeypatch.setenv("AZURE_TENANT_ID", "request-tenant") + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com") + monkeypatch.setenv("AZURE_SCOPE", "https://cognitiveservices.azure.com/.default") + monkeypatch.setattr(azure_common, "azure_ad_cache", DualCache()) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + exchanged = [] + assertions = [] -def _make_anthropic_settings(): - """Settings with ANTHROPIC.KEY configured, no OPENAI.KEY. + def assertion(selector): + assert selector == "oidc/request" + assertions.append(selector) + return "request-assertion" - This simulates the original bug scenario: ANTHROPIC.KEY is set, - but OPENAI.KEY is not, so __init__ falls back to the DUMMY_LITELLM_API_KEY placeholder. - """ - anthropic_key = "test-anthropic-key-12345" - return type("Settings", (), { - "config": type("Config", (), { - "reasoning_effort": None, - "ai_timeout": 30, - "custom_reasoning_model": False, - "max_model_tokens": 32000, - "verbosity_level": 0, - "seed": -1, - "get": lambda self, key, default=None: default, - })(), - "litellm": type("LiteLLM", (), { - "get": lambda self, key, default=None: default, - })(), - "anthropic": type("Anthropic", (), { - "key": anthropic_key - })(), - # Return the Anthropic key when settings.get("ANTHROPIC.KEY") is called - "get": lambda self, key, default=None: ( - anthropic_key if key == "ANTHROPIC.KEY" else default - ), - })() + def exchange(url, *, data): + assert url == "https://login.microsoftonline.com/request-tenant/oauth2/v2.0/token" + assert data["client_id"] == "request-client" + assert data["client_assertion"] == "request-assertion" + assert data["scope"] == "https://cognitiveservices.azure.com/.default" + exchanged.append(data) + return httpx.Response(200, request=httpx.Request("POST", url), json={ + "access_token": "request-ad-token", "expires_in": 3600, + }) + monkeypatch.setattr(azure_common, "get_secret_str", assertion) + monkeypatch.setattr(litellm.module_level_client, "post", exchange) + handler = LiteLLMAIHandler() + second_handler = None + if concurrent_handlers: + monkeypatch.setenv(token_variable, "request-ad-token-second") + second_handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "api_key", "foreign-global-key") + monkeypatch.setenv("AZURE_AD_TOKEN", "later-ad-token") + if alias_mutation == "remove": + monkeypatch.delenv("AZURE_OPENAI_AD_TOKEN", raising=False) + else: + monkeypatch.setenv("AZURE_OPENAI_AD_TOKEN", "later-alias-token") + captured = [] + both_started = asyncio.Event() -class TestApiKeyGuard: + async def send(client, request, **kwargs): + captured.append(request) + if concurrent_handlers: + if len(captured) == 2: + both_started.set() + await both_started.wait() + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) - @pytest.mark.asyncio - async def test_dummy_key_not_forwarded(self, monkeypatch): - """api_key must NOT appear in kwargs when litellm.api_key is the placeholder.""" - monkeypatch.setattr(litellm, "api_key", DUMMY_LITELLM_API_KEY) + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + async def invoke(selected): + if entrypoint == "chat": + await selected.chat_completion("azure/gpt-4o", "sys", "usr") + else: + await selected.probe_completion("azure/gpt-4o") - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() - await handler.chat_completion(model="gpt-4o", system="sys", user="usr") + if concurrent_handlers: + async with asyncio.TaskGroup() as group: + group.create_task(invoke(handler)) + group.create_task(invoke(second_handler)) + else: + for _ in range(request_count): + await invoke(handler) + finally: + await asyncio.sleep(0) + try: + await worker.flush() + finally: + await worker.stop() + assert len(captured) == request_count + if transport == "v1": + expected_url = f"{endpoint}/openai/v1/chat/completions" + elif transport == "cloudflare": + expected_url = f"{endpoint}/gpt-4o/chat/completions?api-version=2024-06-01" + else: + expected_url = f"{endpoint}/openai/deployments/gpt-4o/chat/completions?api-version=2024-06-01" + for index, request in enumerate(captured): + assert str(request.url) == expected_url + authorization = f"Bearer request-ad-token-{index}" if refresh_tokens else expected_authorization + if concurrent_handlers: + assert request.headers.get("Authorization") in {"Bearer request-ad-token", "Bearer request-ad-token-second"} + else: + assert request.headers.get("Authorization") == authorization + assert request.headers.get("api-key") == expected_api_key + assert request.headers["x-request-header"] == "owned" + assert b"request-ad-token" not in request.content + assert b"oidc/request" not in request.content + assert len(exchanged) == (1 if source == "oidc" else 0) + # Native SDK reuse bypasses exchange; the Cloudflare header path resolves + # the assertion per call while retaining the native token exchange cache. + assertion_count = request_count if transport == "cloudflare" else 1 + assert len(assertions) == (assertion_count if source == "oidc" else 0) + if concurrent_handlers: + assert {request.headers.get("Authorization") for request in captured} == { + "Bearer request-ad-token", "Bearer request-ad-token-second", + } + assert litellm.api_key == "foreign-global-key" - assert "api_key" not in mock_call.call_args[1] - @pytest.mark.asyncio - async def test_none_api_key_not_forwarded(self, monkeypatch): - """api_key must NOT appear in kwargs when litellm.api_key is None. +@pytest.mark.parametrize("transport", ("classic", "v1", "cloudflare")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize(("source", "token_variable"), ( + ("token", "AZURE_AD_TOKEN"), + ("token", "AZURE_OPENAI_AD_TOKEN"), + ("oidc", "AZURE_AD_TOKEN"), + ("oidc", "AZURE_OPENAI_AD_TOKEN"), + ("client_credential", None), +)) +@pytest.mark.asyncio +async def test_native_azure_ad_auth_survives_guard_key(monkeypatch, transport, entrypoint, source, token_variable): + await _assert_native_azure_ad_auth(monkeypatch, transport, entrypoint, source, token_variable) - This is the OpenAI-key path: OPENAI.KEY sets litellm.openai_key, - leaving litellm.api_key at None. - """ - monkeypatch.setattr(litellm, "api_key", None) - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() - await handler.chat_completion(model="gpt-4o", system="sys", user="usr") +@pytest.mark.parametrize("transport", ("classic", "v1", "cloudflare")) +@pytest.mark.parametrize("captured_key", ("request-api-key", DUMMY_LITELLM_API_KEY)) +@pytest.mark.parametrize("token_variable", ("AZURE_AD_TOKEN", "AZURE_OPENAI_AD_TOKEN")) +@pytest.mark.asyncio +async def test_native_azure_real_api_key_is_not_a_guard(monkeypatch, transport, captured_key, token_variable): + authorization, api_key = "Bearer request-ad-token", None + if transport == "v1": + authorization = f"Bearer {captured_key}" + elif transport == "cloudflare" and token_variable == "AZURE_AD_TOKEN": + authorization, api_key = None, captured_key + await _assert_native_azure_ad_auth( + monkeypatch, transport, "chat", "token", token_variable, captured_key=captured_key, + expected_authorization=authorization, expected_api_key=api_key, + ) - assert "api_key" not in mock_call.call_args[1] - @pytest.mark.asyncio - async def test_real_key_forwarded(self, monkeypatch): - """api_key IS injected when a real provider key is in litellm.api_key (e.g. Groq, XAI). +@pytest.mark.parametrize("token_variable", ("AZURE_AD_TOKEN", "AZURE_OPENAI_AD_TOKEN")) +@pytest.mark.parametrize("transport", ("classic", "v1", "cloudflare")) +@pytest.mark.parametrize("auth_headers", ( + {"authorization": "Bearer explicit-token"}, + {"API-Key": "explicit-key"}, + {"authorization": "Bearer explicit-token", "API-Key": "explicit-key"}, +)) +@pytest.mark.asyncio +async def test_native_azure_explicit_auth_headers(monkeypatch, token_variable, transport, auth_headers): + authorization = auth_headers.get("authorization") + api_key = auth_headers.get("API-Key") + if authorization is None and (transport != "cloudflare" or token_variable == "AZURE_OPENAI_AD_TOKEN"): + authorization = "Bearer request-ad-token" + await _assert_native_azure_ad_auth( + monkeypatch, transport, "chat", "token", token_variable, auth_headers=auth_headers, + expected_authorization=authorization, expected_api_key=api_key, + ) - The key is set after __init__ to simulate a provider having stored its key there - during initialization, without triggering the placeholder value in __init__. - """ - real_key = "test-provider-key-67890" - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() - # Set after init so __init__'s own dummy-key assignment doesn't overwrite it - monkeypatch.setattr(litellm, "api_key", real_key) - await handler.chat_completion(model="gpt-4o", system="sys", user="usr") +@pytest.mark.parametrize("captured_key", (None, "request-api-key")) +@pytest.mark.parametrize("sdk_alias", (None, "initial-sdk-token")) +@pytest.mark.parametrize("alias_mutation", ("replace", "remove")) +@pytest.mark.parametrize("auth_headers", ( + {}, + {"authorization": "Bearer explicit-token"}, + {"API-Key": "explicit-key"}, + {"authorization": "Bearer explicit-token", "API-Key": "explicit-key"}, +)) +@pytest.mark.asyncio +async def test_native_cloudflare_initial_auth_choice( + monkeypatch, captured_key, sdk_alias, alias_mutation, auth_headers, +): + authorization = auth_headers.get("authorization") + api_key = auth_headers.get("API-Key") + if captured_key is None and api_key is None: + authorization = authorization or "Bearer request-ad-token" + else: + authorization = authorization or (f"Bearer {sdk_alias}" if sdk_alias else None) + api_key = api_key or (None if sdk_alias else captured_key) + await _assert_native_azure_ad_auth( + monkeypatch, "cloudflare", "probe", "token", "AZURE_AD_TOKEN", captured_key=captured_key, + auth_headers=auth_headers, expected_authorization=authorization, expected_api_key=api_key, + sdk_alias=sdk_alias, alias_mutation=alias_mutation, + ) - assert mock_call.call_args[1]["api_key"] == real_key - - @pytest.mark.asyncio - async def test_anthropic_key_not_shadowed_by_dummy_key(self, monkeypatch): - """Original bug scenario: ANTHROPIC.KEY configured without OPENAI.KEY. - - During __init__ the DUMMY_LITELLM_API_KEY placeholder is stored (as the OpenAI - fallback) because OPENAI.KEY is not configured. But litellm.anthropic_key is also - set. The guard must prevent the placeholder from being passed to the call, - allowing litellm to use anthropic_key internally. - - This test replicates the exact bug from GitHub issue #2042. - """ - # Override settings to simulate Anthropic configured, OpenAI not configured - monkeypatch.setattr(litellm_handler, "get_settings", _make_anthropic_settings) - - # Ensure deterministic preconditions: delete OPENAI_API_KEY env var so __init__ - # takes the placeholder branch in litellm_ai_handler.py - monkeypatch.delenv("OPENAI_API_KEY", raising=False) - - # Reset litellm.api_key to avoid cross-test state pollution - monkeypatch.setattr(litellm, "api_key", None) - - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() - - # After init: the placeholder lives in litellm.openai_key (OpenAI fallback) - # and never in litellm.api_key, which LiteLLM checks ahead of provider env - # vars. litellm.anthropic_key holds the real Anthropic key. - assert litellm.openai_key == DUMMY_LITELLM_API_KEY - assert litellm.api_key != DUMMY_LITELLM_API_KEY - - # Call with Anthropic model - await handler.chat_completion( - model="claude-3-5-sonnet-20241022", - system="sys", - user="usr" - ) - # Verify the dummy key was NOT passed to the call. - # This allows litellm to use litellm.anthropic_key internally. - assert "api_key" not in mock_call.call_args[1] - - @pytest.mark.asyncio - async def test_groq_key_forwarded_for_non_ollama_model(self, monkeypatch): - """Regression check for PR #2288: Groq key must be forwarded for non-Ollama models. - - PR #2288 changed the forwarding guard to only forward api_key when - model.startswith('ollama'). This test verifies whether that approach - silently drops the Groq key when calling a non-Ollama model (e.g. gpt-4o). - - Groq sets litellm.api_key during __init__ (see litellm_ai_handler.py line 73) - and relies on it being passed via kwargs["api_key"] to acompletion. - """ - groq_key = "test-groq-key-12345" - - groq_settings = type("Settings", (), { - "config": type("Config", (), { - "reasoning_effort": None, - "ai_timeout": 30, - "custom_reasoning_model": False, - "max_model_tokens": 32000, - "verbosity_level": 0, - "seed": -1, - "get": lambda self, key, default=None: default, - })(), - "litellm": type("LiteLLM", (), { - "get": lambda self, key, default=None: default, - })(), - "groq": type("Groq", (), { - "key": groq_key, - })(), - "get": lambda self, key, default=None: ( - groq_key if key == "GROQ.KEY" else default - ), - })() +@pytest.mark.parametrize("transport", ("classic", "v1", "cloudflare")) +@pytest.mark.parametrize("source", ("client_credential", "oidc")) +@pytest.mark.asyncio +async def test_native_azure_ad_refresh_and_oidc_cache(monkeypatch, transport, source): + await _assert_native_azure_ad_auth( + monkeypatch, transport, "probe", source, "AZURE_AD_TOKEN" if source == "oidc" else None, + request_count=2, refresh_tokens=source == "client_credential", + ) - monkeypatch.setattr(litellm_handler, "get_settings", lambda: groq_settings) - monkeypatch.delenv("OPENAI_API_KEY", raising=False) - monkeypatch.setattr(litellm, "api_key", None) - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() +@pytest.mark.parametrize("transport", ("classic", "v1", "cloudflare")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.asyncio +async def test_native_concurrent_azure_ad_tokens(monkeypatch, transport, entrypoint): + await _assert_native_azure_ad_auth( + monkeypatch, transport, entrypoint, "token", "AZURE_AD_TOKEN", request_count=2, concurrent_handlers=True, + ) - # Confirm __init__ stored the Groq key in litellm.api_key - assert litellm.api_key == groq_key, ( - f"Expected litellm.api_key to be Groq key after __init__, got: {litellm.api_key!r}" - ) - # Call with a non-Ollama model - await handler.chat_completion(model="gpt-4o", system="sys", user="usr") +@pytest.mark.parametrize("model", ("azure_text/gpt-4o", "azure/responses/gpt-4o", "openai/gpt-4o")) +@pytest.mark.asyncio +async def test_azure_header_adapter_does_not_touch_other_transports(model): + handler = LiteLLMAIHandler() + completion = AsyncMock(return_value=_mock_response()) + kwargs = {"model": model, "api_key": DUMMY_LITELLM_API_KEY, "azure_ad_token": "request-token"} + await handler._acompletion(_completion=completion, **kwargs) + assert completion.call_args.kwargs == kwargs - # The Groq key must be forwarded — without it, Groq calls will fail auth - assert mock_call.call_args[1].get("api_key") == groq_key, ( - f"Groq key was NOT forwarded to acompletion. " - f"kwargs had: {mock_call.call_args[1]}" - ) - - @pytest.mark.asyncio - async def test_xai_key_forwarded_for_non_ollama_model(self, monkeypatch): - """Regression check for PR #2288: xAI key must be forwarded for non-Ollama models. - - Similar to Groq, xAI sets litellm.api_key during __init__ and relies on - it being forwarded via kwargs["api_key"]. PR #2288's model-scoped approach - would also break xAI. - """ - xai_key = "xai-test-key-67890" - - xai_settings = type("Settings", (), { - "config": type("Config", (), { - "reasoning_effort": None, - "ai_timeout": 30, - "custom_reasoning_model": False, - "max_model_tokens": 32000, - "verbosity_level": 0, - "seed": -1, - "get": lambda self, key, default=None: default, - })(), - "litellm": type("LiteLLM", (), { - "get": lambda self, key, default=None: default, - })(), - "xai": type("XAI", (), { - "key": xai_key, - })(), - "get": lambda self, key, default=None: ( - xai_key if key == "XAI.KEY" else default - ), - })() - monkeypatch.setattr(litellm_handler, "get_settings", lambda: xai_settings) - monkeypatch.delenv("OPENAI_API_KEY", raising=False) - monkeypatch.setattr(litellm, "api_key", None) +@pytest.mark.asyncio +async def test_azure_oidc_exchange_failure_cannot_use_guard_key(monkeypatch, native_azure_oidc): + state = native_azure_oidc + handler = LiteLLMAIHandler() + exchange = MagicMock(side_effect=RuntimeError("exchange failed")) + monkeypatch.setattr(litellm.module_level_client, "post", exchange) + with pytest.raises(litellm.APIError, match="exchange failed"): + await state.invoke(handler) + exchange.assert_called_once() + assert state.sent == [] + assert litellm.api_key == "foreign-global-key" - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() - assert litellm.api_key == xai_key - await handler.chat_completion(model="gpt-4o", system="sys", user="usr") +@pytest.fixture +async def native_azure_oidc(monkeypatch): + from litellm.caching.caching import DualCache + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.litellm_core_utils import logging_worker + from litellm.llms.azure import common_utils as azure_common - assert mock_call.call_args[1].get("api_key") == xai_key - - @pytest.mark.asyncio - async def test_sambanova_key_forwarded_for_non_ollama_model(self, monkeypatch): - """SambaNova key must be forwarded for non-Ollama models. - - Like Groq and xAI, SambaNova sets litellm.api_key during __init__ - (see litellm_ai_handler.py) and relies on it being forwarded via - kwargs["api_key"] to acompletion. - """ - sambanova_key = "sambanova-test-key-67890" - - sambanova_settings = type("Settings", (), { - "config": type("Config", (), { - "reasoning_effort": None, - "ai_timeout": 30, - "custom_reasoning_model": False, - "max_model_tokens": 32000, - "verbosity_level": 0, - "seed": -1, - "get": lambda self, key, default=None: default, - })(), - "litellm": type("LiteLLM", (), { - "get": lambda self, key, default=None: default, - })(), - "sambanova": type("SambaNova", (), { - "key": sambanova_key, - })(), - "get": lambda self, key, default=None: ( - sambanova_key if key == "SAMBANOVA.KEY" else default - ), - })() + litellm_handler._azure_oidc_entra_provider.cache_clear() + inputs = { + "AZURE_CLIENT_ID": "owned-client", + "AZURE_TENANT_ID": "owned-tenant", + "AZURE_AUTHORITY_HOST": "https://owned-authority.example", + "AZURE_SCOPE": "https://cognitiveservices.azure.com/.default", + } + for name in ("AZURE_CLIENT_SECRET", "AZURE_USERNAME", "AZURE_PASSWORD"): + monkeypatch.delenv(name, raising=False) + for name, value in inputs.items(): + monkeypatch.setenv(name, value) + monkeypatch.setenv("AZURE_AD_TOKEN", "oidc/env/PR_AGENT_TEST_OIDC_ASSERTION") + monkeypatch.setenv("PR_AGENT_TEST_OIDC_ASSERTION", "owned-assertion") + settings = _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": "https://owned.openai.azure.com", + "OPENAI.API_VERSION": "2024-06-01", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + cache = DualCache() + monkeypatch.setattr(azure_common, "azure_ad_cache", cache) + # Preserve warm reuse within a case without inheriting another test's SDK client. + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + native = azure_common.get_azure_ad_token_from_oidc + cache_methods = (cache.get_cache.__func__, cache.set_cache.__func__) + state = SimpleNamespace(inputs=inputs, cache=cache, exchanged=[], sent=[], status=200, ttl=3600, ambient=None) + state.expected_api_key = None + state.expected_url = ( + "https://owned.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-06-01" + ) + state.response_kind = "chat" + state.stream = False - monkeypatch.setattr(litellm_handler, "get_settings", lambda: sambanova_settings) - monkeypatch.delenv("OPENAI_API_KEY", raising=False) - monkeypatch.setattr(litellm, "api_key", None) + def exchange(url, *, data): + if state.ambient is not None: + assert dict(os.environ) == state.ambient + assert azure_common.azure_ad_cache is cache + assert (cache.get_cache.__func__, cache.set_cache.__func__) == cache_methods + state.exchanged.append((url, dict(data))) + return httpx.Response(state.status, request=httpx.Request("POST", url), json={ + "access_token": f"exchanged-token-{len(state.exchanged)}", "expires_in": state.ttl, + }) - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() + async def send(client, request, **kwargs): + state.sent.append(request) + expected_keys = [] if state.expected_api_key is None else [state.expected_api_key] + assert request.headers.get_list("api-key") == expected_keys + assert str(request.url) == state.expected_url + assert "oidc/" not in str(request.headers) + assert "assertion" not in str(request.headers) + assert b"oidc/" not in request.content + assert b"assertion" not in request.content + assert b"exchanged-token" not in request.content + if state.response_kind == "responses": + return httpx.Response(200, request=request, json={ + "id": "resp_test", "object": "response", "created_at": 0, "model": "gpt-4o", + "status": "completed", "error": None, "incomplete_details": None, + "output": [{"id": "msg_test", "type": "message", "role": "assistant", "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}]}], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }) + if state.response_kind == "text": + payload = { + "id": "test", "object": "text_completion", "created": 0, "model": "gpt-3.5-turbo-instruct", + "choices": [{"index": 0, "text": "ok", "logprobs": None, "finish_reason": "stop"}], + } + if state.stream: + return httpx.Response(200, request=request, headers={"content-type": "text/event-stream"}, + content=f"data: {json.dumps(payload)}\n\ndata: [DONE]\n\n".encode()) + return httpx.Response(200, request=request, json=payload) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) - assert litellm.api_key == sambanova_key - await handler.chat_completion( - model="sambanova/MiniMax-M3", system="sys", user="usr" - ) + async def invoke(handler, entrypoint="probe", guard_key=True, model="azure/gpt-4o"): + # Cover native keyless dispatch as well as the request-local guard used + # to prevent a process-wide API key from replacing OIDC authentication. + ambient_key = "foreign-global-key" if guard_key else None + monkeypatch.setattr(litellm, "api_key", ambient_key) + environment = dict(os.environ) + state.ambient = environment + try: + if entrypoint == "chat": + await handler.chat_completion(model, "sys", "usr") + else: + await handler.probe_completion(model) + finally: + assert dict(os.environ) == environment + assert azure_common.azure_ad_cache is cache + assert (cache.get_cache.__func__, cache.set_cache.__func__) == cache_methods + assert litellm.api_key == ambient_key + assert litellm_handler._azure_oidc_request.get() is None + assert litellm_handler._azure_ad_responses_request.get() is None + return state.sent[-1].headers.get("Authorization") - assert mock_call.call_args[1].get("api_key") == sambanova_key - - @pytest.mark.asyncio - async def test_databricks_model_does_not_forward_foreign_key(self, monkeypatch): - """Databricks models authenticate via DATABRICKS_API_KEY/DATABRICKS_API_BASE env vars. - - In a multi-provider config another provider (e.g. Groq/OpenRouter) may have stored - its key in litellm.api_key during __init__. That key must NOT be forwarded for - databricks/* calls, otherwise it would override the intended env-var auth and break - Databricks authentication. - """ - foreign_key = "test-groq-key-shadowing-databricks" - - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() - # Simulate another provider having populated litellm.api_key during init - monkeypatch.setattr(litellm, "api_key", foreign_key) - await handler.chat_completion( - model="databricks/databricks-claude-sonnet-4", system="sys", user="usr" - ) + state.invoke = invoke + monkeypatch.setattr(litellm.module_level_client, "post", exchange) + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + yield state + finally: + litellm_handler._azure_oidc_entra_provider.cache_clear() + assert native.__globals__["os"] is os + await asyncio.sleep(0) + try: + await worker.flush() + finally: + await worker.stop() - assert "api_key" not in mock_call.call_args[1], ( - f"Foreign provider key must not be forwarded for databricks/* models. " - f"kwargs had: {mock_call.call_args[1]}" - ) - - @pytest.mark.asyncio - async def test_databricks_model_does_not_forward_foreign_api_base(self, monkeypatch): - """Databricks models select their endpoint via the DATABRICKS_API_BASE env var. - - In a multi-provider config another provider (OpenRouter/Ollama/Azure AD/OpenAI) may - have set self.api_base during __init__. That base URL must NOT be forwarded for - databricks/* calls, otherwise it would route the request to the wrong host and - override the intended DATABRICKS_API_BASE endpoint. The Databricks base (or None, - which lets LiteLLM read the env var) must be used instead. - """ - databricks_base = "https://adb-1234.azuredatabricks.net/serving-endpoints" - monkeypatch.setenv("DATABRICKS_API_BASE", databricks_base) - - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() - # Simulate another provider having set api_base during init - handler.api_base = "https://openrouter.ai/api/v1" - await handler.chat_completion( - model="databricks/databricks-claude-sonnet-4", system="sys", user="usr" - ) - assert mock_call.call_args[1]["api_base"] == databricks_base, ( - f"Databricks endpoint must come from DATABRICKS_API_BASE, not a foreign provider's " - f"api_base. kwargs had: {mock_call.call_args[1]}" - ) - - @pytest.mark.asyncio - async def test_databricks_guards_survive_azure_mode(self, monkeypatch): - """Azure mode must not rewrite databricks/* models and bypass the Databricks guards. - - When Azure is enabled in a multi-provider config (OPENAI.API_TYPE=azure or AZURE_AD), - chat_completion() prepends 'azure/' to the model. If that rewrite happened for a - databricks/* model the prefix-based guards would never trigger, routing the call to - Azure with a foreign key/base. The model must keep its 'databricks/' prefix, the foreign - key must not be forwarded, and api_base must come from DATABRICKS_API_BASE. - """ - foreign_key = "test-azure-key-shadowing-databricks" - databricks_base = "https://adb-1234.azuredatabricks.net/serving-endpoints" - monkeypatch.setenv("DATABRICKS_API_BASE", databricks_base) - - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() - # Simulate Azure mode + a foreign key/base set by another provider during init - handler.azure = True - handler.api_base = "https://my-azure.openai.azure.com" - monkeypatch.setattr(litellm, "api_key", foreign_key) - await handler.chat_completion( - model="databricks/databricks-claude-sonnet-4", system="sys", user="usr" - ) +@pytest.mark.parametrize("field", ("AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_AUTHORITY_HOST", "AZURE_SCOPE")) +@pytest.mark.parametrize("mutation", ("changed", "removed", "absent", "empty")) +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_oidc_outer_snapshot(monkeypatch, native_azure_oidc, field, mutation, guard_key): + state = native_azure_oidc + captured = dict(state.inputs) + if field == "AZURE_SCOPE": + captured[field] = "owned-scope" + monkeypatch.setenv(field, captured[field]) + if mutation == "absent": + monkeypatch.delenv(field) + captured[field] = None + elif mutation == "empty": + monkeypatch.setenv(field, "") + captured[field] = "" + handler = LiteLLMAIHandler() + if mutation == "removed": + monkeypatch.delenv(field) + else: + monkeypatch.setenv(field, "https://foreign.example" if "HOST" in field else "foreign-value") + if mutation == "absent" and field in {"AZURE_CLIENT_ID", "AZURE_TENANT_ID"}: + with pytest.raises(litellm.BadRequestError, match="AZURE_CLIENT_ID and AZURE_TENANT_ID must be set"): + await state.invoke(handler, guard_key=guard_key) + assert state.exchanged == state.sent == [] + return + assert await state.invoke(handler, guard_key=guard_key) == "Bearer exchanged-token-1" + authority = captured["AZURE_AUTHORITY_HOST"] + if authority is None: + authority = "https://login.microsoftonline.com" + scope = captured["AZURE_SCOPE"] + if scope is None: + scope = "https://cognitiveservices.azure.com/.default" + assert state.exchanged == [(f"{authority}/{captured['AZURE_TENANT_ID']}/oauth2/v2.0/token", { + "client_id": captured["AZURE_CLIENT_ID"], "grant_type": "client_credentials", "scope": scope, + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": "owned-assertion", + })] - forwarded = mock_call.call_args[1] - assert forwarded["model"] == "databricks/databricks-claude-sonnet-4", ( - f"databricks/* model must not be rewritten with an 'azure/' prefix. " - f"kwargs had: {forwarded}" - ) - assert "api_key" not in forwarded, ( - f"Foreign provider key must not be forwarded for databricks/* models even in Azure " - f"mode. kwargs had: {forwarded}" - ) - assert forwarded["api_base"] == databricks_base, ( - f"Databricks endpoint must come from DATABRICKS_API_BASE even in Azure mode. " - f"kwargs had: {forwarded}" - ) - - @pytest.mark.asyncio - async def test_ollama_and_groq_coexist(self, monkeypatch): - """Verify both Ollama and Groq keys can coexist and be forwarded correctly. - - When multiple providers are configured, litellm.api_key gets overwritten - sequentially during __init__. The sentinel guard should still forward - whatever real key is currently in litellm.api_key. - """ - groq_key = "gsk-groq-key" - ollama_key = "ollama-key" - - # Simulate: Groq key set first, then Ollama overwrites litellm.api_key - mixed_settings = type("Settings", (), { - "config": type("Config", (), { - "reasoning_effort": None, - "ai_timeout": 30, - "custom_reasoning_model": False, - "max_model_tokens": 32000, - "verbosity_level": 0, - "seed": -1, - "get": lambda self, key, default=None: default, - })(), - "litellm": type("LiteLLM", (), { - "get": lambda self, key, default=None: default, - })(), - "groq": type("Groq", (), {"key": groq_key})(), - "ollama": type("Ollama", (), { - "api_key": ollama_key, - "api_base": "http://localhost:11434", - })(), - "get": lambda self, key, default=None: ( - groq_key if key == "GROQ.KEY" else - ollama_key if key == "OLLAMA.API_KEY" else - "http://localhost:11434" if key == "OLLAMA.API_BASE" else - default - ), - })() - monkeypatch.setattr(litellm_handler, "get_settings", lambda: mixed_settings) - monkeypatch.delenv("OPENAI_API_KEY", raising=False) - monkeypatch.setattr(litellm, "api_key", None) +@pytest.mark.asyncio +async def test_native_azure_oidc_thread_boundary_snapshot(monkeypatch, native_azure_oidc): + state = native_azure_oidc + handler = LiteLLMAIHandler() + loop = asyncio.get_running_loop() + run_in_executor = loop.run_in_executor + boundaries = [] - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() + def mutate_then_dispatch(executor, func, *args): + # LiteLLM dispatches native completion through this executor boundary. + if not boundaries: + boundaries.append(func) + for name in state.inputs: + monkeypatch.setenv(name, "https://foreign.example" if "HOST" in name else "foreign-value") + state.ambient = dict(os.environ) + return run_in_executor(executor, func, *args) - # After init, litellm.api_key should be Ollama (last assignment) - assert litellm.api_key == ollama_key + monkeypatch.setattr(loop, "run_in_executor", mutate_then_dispatch) + monkeypatch.setattr(litellm, "api_key", "foreign-global-key") + # This test deliberately changes the environment at dispatch, outside the adapter. + await handler.chat_completion("azure/gpt-4o", "sys", "usr") + assert len(boundaries) == 1 + url, data = state.exchanged[0] + assert url == "https://owned-authority.example/owned-tenant/oauth2/v2.0/token" + assert data["client_id"] == "owned-client" + assert data["scope"] == state.inputs["AZURE_SCOPE"] + assert state.sent[0].headers["Authorization"] == "Bearer exchanged-token-1" + assert all("foreign" in os.environ[name] for name in state.inputs) - # Call with Ollama model — should get Ollama key - await handler.chat_completion(model="ollama/mistral", system="sys", user="usr") - assert mock_call.call_args[1]["api_key"] == ollama_key - # Call with non-Ollama model — should still forward the key - # (which is Ollama in this case, but the guard correctly allows real keys through) - await handler.chat_completion(model="gpt-4o", system="sys", user="usr") - assert mock_call.call_args[1]["api_key"] == ollama_key - - -class _StopCall(Exception): - """Sentinel raised from a patched LiteLLM transport once the api_key is captured.""" - - -class TestPlaceholderDoesNotShadowProviderEnvVars: - """Regression tests for issue #2544. - - The placeholder must not sit in the global ``litellm.api_key``: LiteLLM resolves - several providers as ``api_key or litellm.api_key or or - get_secret("_API_KEY")``, so a truthy placeholder there wins over the - user's provider env var (OpenRouter, Azure, Mistral, ... ) and the request goes - out with "dummy_key". Keeping the placeholder in ``litellm.openai_key`` — which - LiteLLM only consults on the OpenAI/OpenAI-compatible paths — preserves the - keyless local-endpoint use case without shadowing anything else. - """ - - @pytest.fixture(autouse=True) - def restore_openai_key(self, monkeypatch): - monkeypatch.setattr(litellm, "openai_key", None) - monkeypatch.setattr(litellm, "api_key", None) - monkeypatch.delenv("OPENAI_API_KEY", raising=False) - - @pytest.mark.asyncio - async def test_keyless_init_clears_a_stale_provider_key(self, monkeypatch): - """A keyless request must not inherit the previous request's provider key. - - __init__ mutates process-global LiteLLM state, so a provider branch from an - earlier request (Groq/xAI/SambaNova/OpenRouter all write litellm.api_key) can - still be sitting there. chat_completion() forwards any truthy non-placeholder - litellm.api_key, so without an explicit reset a keyless request would - authenticate with — and leak — the earlier request's credential. - """ - groq_settings = type("Settings", (), { - "config": type("Config", (), { - "reasoning_effort": None, - "ai_timeout": 30, - "custom_reasoning_model": False, - "max_model_tokens": 32000, - "verbosity_level": 0, - "seed": -1, - "get": lambda self, key, default=None: default, - })(), - "litellm": type("LiteLLM", (), { - "get": lambda self, key, default=None: default, - })(), - "groq": type("Groq", (), {"key": "gsk-tenant-a"})(), - "get": lambda self, key, default=None: ( - "gsk-tenant-a" if key == "GROQ.KEY" else default - ), - })() +@pytest.mark.parametrize("transport", ("classic", "v1", "responses", "text", "cloudflare", "cloudflare_text")) +@pytest.mark.parametrize("token_kind", ("ordinary", "oidc")) +@pytest.mark.parametrize("key", (None, "owned-key", "dummy_key")) +@pytest.mark.parametrize("late_key", (None, "api_key", "azure_key", "AZURE_API_KEY", "AZURE_OPENAI_API_KEY")) +@pytest.mark.asyncio +async def test_native_azure_token_executor_boundary( + monkeypatch, native_azure_oidc, transport, token_kind, key, late_key, +): + state = native_azure_oidc + if token_kind == "ordinary": + monkeypatch.setenv("AZURE_AD_TOKEN", "owned-ad-token") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + endpoint = "https://owned.openai.azure.com" + if transport.startswith("cloudflare"): + endpoint = "https://gateway.ai.cloudflare.com/v1/account/gateway/azure-openai/resource" + version = "v1" if transport == "v1" else "2024-06-01" + settings = _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": endpoint, + "OPENAI.API_VERSION": version, "OPENAI.KEY": key, + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + model = "azure/gpt-4o" + if transport in ("text", "cloudflare_text"): + model = "azure_text/gpt-3.5-turbo-instruct" + state.response_kind = "text" + elif transport == "responses": + model = "azure/responses/gpt-4o" + state.response_kind = "responses" + if transport == "v1": + state.expected_url = f"{endpoint}/openai/v1/chat/completions" + elif transport == "responses": + state.expected_url = f"{endpoint}/openai/responses?api-version={version}" + elif transport.startswith("cloudflare"): + suffix = "gpt-3.5-turbo-instruct/completions" if transport.endswith("text") else "gpt-4o/chat/completions" + state.expected_url = f"{endpoint}/{suffix}?api-version={version}" + elif transport == "text": + state.expected_url = f"{endpoint}/openai/deployments/gpt-3.5-turbo-instruct/completions?api-version={version}" + expected = "Bearer owned-ad-token" if token_kind == "ordinary" else "Bearer exchanged-token-1" + if key and (transport == "responses" or transport.startswith("cloudflare")): + expected, state.expected_api_key = None, key + elif key and transport == "v1": + expected = f"Bearer {key}" + handler = LiteLLMAIHandler() + loop = asyncio.get_running_loop() + run_in_executor = loop.run_in_executor + boundaries = [] + + def mutate_then_dispatch(executor, function, *args): + if not boundaries: + boundaries.append(function) + if late_key in ("api_key", "azure_key"): + monkeypatch.setattr(litellm, late_key, "foreign-key") + elif late_key: + monkeypatch.setenv(late_key, "foreign-key") + state.ambient = dict(os.environ) + return run_in_executor(executor, function, *args) + + monkeypatch.setattr(loop, "run_in_executor", mutate_then_dispatch) + await handler.probe_completion(model) + assert boundaries + assert len(state.sent) == 1 + assert state.sent[0].headers.get("Authorization") == expected + assert litellm_handler._azure_oidc_request.get() is None + assert litellm_handler._azure_ad_responses_request.get() is None + + +@pytest.mark.parametrize("identity", ("same", "client", "tenant", "authority", "scope", "default", "empty", "legacy")) +@pytest.mark.asyncio +async def test_native_azure_oidc_scope_cache(monkeypatch, native_azure_oidc, identity): + state = native_azure_oidc + if identity == "default": + monkeypatch.delenv("AZURE_SCOPE") + native_key = json.dumps({ + "azure_client_id": "owned-client", "azure_tenant_id": "owned-tenant", + "azure_authority_host": "https://owned-authority.example", "oidc_token": "owned-assertion", + }) + if identity == "legacy": + state.cache.set_cache(native_key, "old-unnamespaced-token", ttl=3600) + first = LiteLLMAIHandler() + assert await state.invoke(first) == "Bearer exchanged-token-1" + fields = {"client": "AZURE_CLIENT_ID", "tenant": "AZURE_TENANT_ID", "authority": "AZURE_AUTHORITY_HOST"} + if identity in fields: + monkeypatch.setenv(fields[identity], "https://second.example" if identity == "authority" else "second") + elif identity in {"scope", "empty", "default"}: + monkeypatch.setenv("AZURE_SCOPE", { + "scope": "second-scope", "empty": "", "default": state.inputs["AZURE_SCOPE"], + }[identity]) + second = LiteLLMAIHandler() + count = 1 if identity in {"same", "default", "legacy"} else 2 + assert await state.invoke(second) == f"Bearer exchanged-token-{count}" + assert len(state.exchanged) == count + # An older handler keeps its identity even after a second handler is captured. + assert await state.invoke(first) == "Bearer exchanged-token-1" + assert len(state.exchanged) == count + namespaced_key = json.dumps(["pr-agent.azure-oidc.v1", state.inputs["AZURE_SCOPE"], native_key]) + assert state.cache.get_cache(namespaced_key) == "exchanged-token-1" + assert state.cache.get_cache(native_key) == ("old-unnamespaced-token" if identity == "legacy" else None) + + +@pytest.mark.parametrize("refresh", ("unchanged", "assertion", "expiry")) +@pytest.mark.asyncio +async def test_native_azure_oidc_assertion_rotation(monkeypatch, native_azure_oidc, refresh): + state = native_azure_oidc + # Responses executes token selection per request, so this exercises the + # token cache independently of native SDK client reuse (covered below). + state.response_kind = "responses" + state.expected_url = "https://owned.openai.azure.com/openai/responses?api-version=2024-06-01" + model = "azure/responses/gpt-4o" + if refresh == "expiry": + state.ttl = -1 + handler = LiteLLMAIHandler() + assert await state.invoke(handler, "chat", model=model) == "Bearer exchanged-token-1" + monkeypatch.setenv("AZURE_AD_TOKEN", "oidc/env/PR_AGENT_TEST_FOREIGN_ASSERTION") + monkeypatch.setenv("PR_AGENT_TEST_FOREIGN_ASSERTION", "foreign-assertion") + if refresh == "assertion": + monkeypatch.setenv("PR_AGENT_TEST_OIDC_ASSERTION", "rotated-assertion") + count = 1 if refresh == "unchanged" else 2 + assert await state.invoke(handler, "chat", model=model) == f"Bearer exchanged-token-{count}" + expected = ["owned-assertion"] + if count == 2: + expected.append("rotated-assertion" if refresh == "assertion" else "owned-assertion") + assert [data["client_assertion"] for _, data in state.exchanged] == expected + + +@pytest.mark.parametrize("failure", ( + "assertion", "exchange", "interface", "interface_kind", "interface_default", "snapshot", +)) +@pytest.mark.asyncio +async def test_native_azure_oidc_fail_closed(monkeypatch, native_azure_oidc, failure): + from litellm.llms.azure import common_utils as azure_common + + state = native_azure_oidc + handler = LiteLLMAIHandler() + if failure == "assertion": + monkeypatch.delenv("PR_AGENT_TEST_OIDC_ASSERTION") + error, message = litellm.APIError, "PR_AGENT_TEST_OIDC_ASSERTION not found" + elif failure == "exchange": + state.status = 401 + error, message = litellm.AuthenticationError, "exchanged-token" + elif failure == "snapshot": + del handler._azure_oidc_environment["AZURE_SCOPE"] + error, message = RuntimeError, "Azure OIDC exchange snapshot is incomplete" + elif failure in {"interface_kind", "interface_default"}: + # Mutate the native interface, not a previously installed bridge wrapper. + native = inspect.unwrap(azure_common.get_azure_ad_token_from_oidc) + signature = inspect.signature(native) + parameters = list(signature.parameters.values()) + parameter = parameters[-1] + parameters[-1] = ( + parameter.replace(kind=inspect.Parameter.KEYWORD_ONLY) + if failure == "interface_kind" else parameter.replace(default="changed-default") + ) + monkeypatch.setattr(native, "__signature__", signature.replace(parameters=parameters), raising=False) + error, message = litellm.APIError, "Azure OIDC exchange interface is incompatible" + else: + # A callable with an incompatible native signature must never be invoked. + monkeypatch.setattr(azure_common, "get_azure_ad_token_from_oidc", lambda selector: "unsafe-token") + error, message = RuntimeError, "(?i)azure|oidc|incompatible" + with pytest.raises(error, match=message): + await state.invoke(handler) + assert len(state.exchanged) == (1 if failure == "exchange" else 0) + assert state.sent == [] + + +@pytest.mark.parametrize("mutation", ("changed", "removed")) +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_oidc_bridge_responses(monkeypatch, native_azure_oidc, mutation, guard_key): + from litellm.llms.azure import common_utils as azure_common + + state = native_azure_oidc + state.response_kind = "responses" + state.expected_url = "https://owned.openai.azure.com/openai/responses?api-version=2024-06-01" + handler = LiteLLMAIHandler() + for field in ("AZURE_CLIENT_ID", "AZURE_TENANT_ID"): + if mutation == "removed": + monkeypatch.delenv(field) + else: + monkeypatch.setenv(field, "foreign-identity") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + provider = MagicMock(side_effect=AssertionError("OIDC must not create an additional credential provider")) + monkeypatch.setattr(azure_common, "get_azure_ad_token_provider", provider) + assert await state.invoke(handler, guard_key=guard_key, model="azure/responses/gpt-4o") == "Bearer exchanged-token-1" + assert len(state.exchanged) == 1 + url, data = state.exchanged[0] + assert url == "https://owned-authority.example/owned-tenant/oauth2/v2.0/token" + assert data["client_id"] == "owned-client" + provider.assert_not_called() + + +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("key", (None, "owned-key", "dummy_key")) +@pytest.mark.parametrize("ambient", ("none", "global", "azure", "environment")) +@pytest.mark.asyncio +async def test_native_azure_ad_responses_guard(monkeypatch, native_azure_oidc, entrypoint, key, ambient): + state = native_azure_oidc + state.response_kind = "responses" + state.expected_url = "https://owned.openai.azure.com/openai/responses?api-version=2024-06-01" + state.expected_api_key = key + monkeypatch.setenv("AZURE_AD_TOKEN", "owned-ad-token") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + settings = _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": "https://owned.openai.azure.com", + "OPENAI.API_VERSION": "2024-06-01", "OPENAI.KEY": key, + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + handler = LiteLLMAIHandler() + if ambient == "azure": + monkeypatch.setattr(litellm, "azure_key", "foreign-key") + elif ambient == "environment": + monkeypatch.setenv("AZURE_API_KEY", "foreign-key") + token = await state.invoke(handler, entrypoint=entrypoint, guard_key=ambient == "global", + model="azure/responses/gpt-4o") + assert token == (None if key else "Bearer owned-ad-token") + assert len(state.sent) == 1 + assert state.exchanged == [] + + +@pytest.mark.parametrize("source", ("client_secret", "password")) +@pytest.mark.parametrize("timing", ("initial", "late")) +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.parametrize("refresh", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_ad_responses_companion(monkeypatch, native_azure_oidc, source, timing, guard_key, refresh): + import azure.identity - monkeypatch.setattr(litellm_handler, "get_settings", lambda: groq_settings) - LiteLLMAIHandler() # request 1: tenant A, Groq key lands in litellm.api_key - assert litellm.api_key == "gsk-tenant-a" - - monkeypatch.setattr(litellm_handler, "get_settings", _make_settings) - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_response() - handler = LiteLLMAIHandler() # request 2: tenant B, no keys configured - await handler.chat_completion( - model="openai/local-model", system="sys", user="usr" + state = native_azure_oidc + state.response_kind = "responses" + state.expected_url = "https://owned.openai.azure.com/openai/responses?api-version=2024-06-01" + monkeypatch.setenv("AZURE_AD_TOKEN", "owned-ad-token") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", refresh) + fields = {"AZURE_CLIENT_SECRET": "owned-secret"} if source == "client_secret" else { + "AZURE_USERNAME": "owned-user", "AZURE_PASSWORD": "owned-password", + } + if timing == "initial": + for name, value in fields.items(): + monkeypatch.setenv(name, value) + handler = LiteLLMAIHandler() + for name in fields: + monkeypatch.setenv(name, "foreign-value") + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://foreign-authority.example") + credential_name = "ClientSecretCredential" if source == "client_secret" else "UsernamePasswordCredential" + constructor = MagicMock(wraps=getattr(azure.identity, credential_name)) + monkeypatch.setattr(azure.identity, credential_name, constructor) + provider = MagicMock(return_value=lambda: "provider-token") + monkeypatch.setattr(azure.identity, "get_bearer_token_provider", provider) + if refresh and timing == "late": + with pytest.raises(litellm.APIConnectionError, match="implicit Azure AD credential discovery"): + await state.invoke(handler, model="azure/responses/gpt-4o", guard_key=guard_key) + constructor.assert_not_called() + provider.assert_not_called() + assert state.sent == state.exchanged == [] + return + expected = "provider-token" if timing == "initial" else "owned-ad-token" + assert await state.invoke(handler, model="azure/responses/gpt-4o", guard_key=guard_key) == f"Bearer {expected}" + assert state.exchanged == [] + if timing == "initial": + constructor.assert_called_once() + provider.assert_called_once() + assert constructor.call_args.kwargs["authority"] == "https://owned-authority.example" + if source == "client_secret": + assert constructor.call_args.args == ("owned-tenant", "owned-client", "owned-secret") + else: + assert constructor.call_args.kwargs["username"] == "owned-user" + assert constructor.call_args.kwargs["password"] == "owned-password" + else: + constructor.assert_not_called() + provider.assert_not_called() + + +@pytest.mark.parametrize("transport", ("classic", "v1")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("source", ("client_secret", "password", "both", "incomplete", "password_without_tenant")) +@pytest.mark.parametrize("timing", ("initial", "late")) +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.parametrize("key", (None, "owned-key", "dummy_key")) +@pytest.mark.parametrize("refresh", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_ad_sdk_companion( + monkeypatch, native_azure_oidc, transport, entrypoint, source, timing, guard_key, key, refresh, +): + import azure.identity + + state = native_azure_oidc + monkeypatch.setenv("AZURE_AD_TOKEN", "owned-ad-token") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", refresh) + version = "v1" if transport == "v1" else "2024-06-01" + settings = _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": "https://owned.openai.azure.com", + "OPENAI.API_VERSION": version, "OPENAI.KEY": key, + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + if transport == "v1": + state.expected_url = "https://owned.openai.azure.com/openai/v1/chat/completions" + fields = {} + if source in ("client_secret", "both", "incomplete", "password_without_tenant"): + fields["AZURE_CLIENT_SECRET"] = "owned-secret" + if source in ("password", "both", "password_without_tenant"): + fields.update(AZURE_USERNAME="owned-user", AZURE_PASSWORD="owned-password") + if source in ("incomplete", "password_without_tenant"): + monkeypatch.delenv("AZURE_TENANT_ID") + if timing == "initial": + for name, value in fields.items(): + monkeypatch.setenv(name, value) + handler = LiteLLMAIHandler() + for name in fields: + monkeypatch.setenv(name, "foreign-value") + monkeypatch.setenv("AZURE_TENANT_ID", "foreign-tenant") + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://foreign-authority.example") + constructors = {} + for name in ("ClientSecretCredential", "UsernamePasswordCredential"): + constructors[name] = MagicMock(wraps=getattr(azure.identity, name)) + monkeypatch.setattr(azure.identity, name, constructors[name]) + token_provider = MagicMock(return_value="companion-token") + provider_factory = MagicMock(return_value=token_provider) + monkeypatch.setattr(azure.identity, "get_bearer_token_provider", provider_factory) + if refresh and not key and (timing == "late" or source == "incomplete"): + with pytest.raises(Exception, match="implicit Azure AD credential discovery"): + await state.invoke(handler, entrypoint, guard_key=guard_key) + for constructor in constructors.values(): + constructor.assert_not_called() + provider_factory.assert_not_called() + assert state.sent == [] + assert state.exchanged == [] + return + expected = "owned-ad-token" + if transport == "v1": + expected = key or ("companion-token" if timing == "initial" and source != "incomplete" else expected) + assert await state.invoke(handler, entrypoint, guard_key=guard_key) == f"Bearer {expected}" + if expected == "companion-token": + token_provider.assert_called() + selected = None + if timing == "initial": + if source in ("client_secret", "both") and not key: + selected = "ClientSecretCredential" + elif source in ("password", "both", "password_without_tenant"): + selected = "UsernamePasswordCredential" + for name, constructor in constructors.items(): + if name != selected: + constructor.assert_not_called() + continue + constructor.assert_called_once() + assert constructor.call_args.kwargs["authority"] == "https://owned-authority.example" + if name == "ClientSecretCredential": + assert constructor.call_args.args == ("owned-tenant", "owned-client", "owned-secret") + else: + assert constructor.call_args.kwargs["username"] == "owned-user" + assert constructor.call_args.kwargs["password"] == "owned-password" + assert state.exchanged == [] + + +@pytest.mark.parametrize("transport", ("classic", "v1")) +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.asyncio +async def test_native_azure_ad_sdk_rejects_implicit_refresh(monkeypatch, native_azure_oidc, transport, guard_key, entrypoint): + import azure.identity + + state = native_azure_oidc + monkeypatch.setenv("AZURE_AD_TOKEN", "owned-ad-token") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + settings = _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": "https://owned.openai.azure.com", + "OPENAI.API_VERSION": "v1" if transport == "v1" else "2024-06-01", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + if transport == "v1": + state.expected_url = "https://owned.openai.azure.com/openai/v1/chat/completions" + handler = LiteLLMAIHandler() + for name in ("AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_CLIENT_SECRET"): + monkeypatch.setenv(name, "foreign-identity") + constructor = MagicMock() + provider_factory = MagicMock(return_value=lambda: "foreign-token") + monkeypatch.setattr(azure.identity, "ClientSecretCredential", constructor) + monkeypatch.setattr(azure.identity, "get_bearer_token_provider", provider_factory) + with pytest.raises(Exception, match="implicit Azure AD credential discovery"): + await state.invoke(handler, entrypoint, guard_key=guard_key) + constructor.assert_not_called() + provider_factory.assert_not_called() + assert state.sent == [] + assert state.exchanged == [] + + +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_ad_sdk_reuses_safe_client_when_refresh_enabled(monkeypatch, native_azure_oidc, guard_key): + from litellm.llms.azure import common_utils as azure_common + + state = native_azure_oidc + monkeypatch.setenv("AZURE_AD_TOKEN", "owned-ad-token") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + settings = _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": "https://owned.openai.azure.com", "OPENAI.API_VERSION": "v1", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + state.expected_url = "https://owned.openai.azure.com/openai/v1/chat/completions" + handler = LiteLLMAIHandler() + assert await state.invoke(handler, guard_key=guard_key) == "Bearer owned-ad-token" + factory = MagicMock(side_effect=AssertionError("Implicit discovery must not run for a cached client")) + monkeypatch.setattr(azure_common, "get_azure_ad_token_provider", factory) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + monkeypatch.setenv("AZURE_CLIENT_SECRET", "foreign-secret") + assert await state.invoke(handler, guard_key=guard_key) == "Bearer owned-ad-token" + factory.assert_not_called() + assert len(state.sent) == 2 + + +@pytest.fixture +async def native_companion_auth(monkeypatch, native_azure_oidc): + from azure.core.pipeline.transport import RequestsTransport + from litellm.llms.azure import common_utils + + caches = [litellm.in_memory_llm_clients_cache] + common_utils._cached_entra_id_token_provider.cache_clear() + + def deny_auth_http(*args, **kwargs): + raise AssertionError("Unexpected authentication HTTP") + + monkeypatch.setattr(RequestsTransport, "send", deny_auth_http) + monkeypatch.setattr(httpx.Client, "send", deny_auth_http) + try: + yield caches + finally: + common_utils._cached_entra_id_token_provider.cache_clear() + clients = {id(value): value for cache in caches for value in cache.cache_dict.values()} + for client in clients.values(): + result = client.close() + if inspect.isawaitable(result): + await result + + +@pytest.mark.parametrize("transport", ( + "classic", "v1", "responses", "text", "raw", "raw_key_host", "sdk_alias", "responses_alias", +)) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("source", ("client_secret", "password")) +@pytest.mark.parametrize("key", (None, "owned-key", "dummy_key")) +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.asyncio +async def test_native_tokenless_azure_companion( + monkeypatch, native_azure_oidc, native_companion_auth, transport, entrypoint, source, key, guard_key, +): + import azure.identity + from litellm.caching.llm_caching_handler import LLMClientCache + + state = native_azure_oidc + monkeypatch.delenv("AZURE_AD_TOKEN") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + if source == "client_secret": + monkeypatch.setenv("AZURE_CLIENT_SECRET", "owned-secret") + else: + monkeypatch.setenv("AZURE_USERNAME", "owned-user") + monkeypatch.setenv("AZURE_PASSWORD", "owned-password") + version = "v1" if transport == "v1" else "2024-06-01" + endpoint = "https://owned.openai.azure.com" + settings = _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": endpoint, + "OPENAI.API_VERSION": version, "OPENAI.KEY": key, + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + model = "azure/gpt-4o" + if transport == "v1": + state.expected_url = f"{endpoint}/openai/v1/chat/completions" + elif transport == "responses": + model = "azure/responses/gpt-4o" + state.response_kind = "responses" + state.expected_url = f"{endpoint}/openai/responses?api-version={version}" + elif transport == "text": + model = "azure_text/gpt-3.5-turbo-instruct" + state.response_kind = "text" + state.expected_url = f"{endpoint}/openai/deployments/gpt-3.5-turbo-instruct/completions?api-version={version}" + elif transport.startswith("raw"): + # OpenAI model names on azure_ai normalize to the Azure SDK instead. + model = "azure_ai/test-model" + endpoint = "https://owned.services.ai.azure.com" if transport == "raw_key_host" else "https://owned.example" + monkeypatch.setenv("AZURE_AI_API_BASE", endpoint) + if key: + monkeypatch.setenv("AZURE_AI_API_KEY", key) + settings = _make_settings() + suffix = "/models/chat/completions" if transport == "raw_key_host" else "/chat/completions" + state.expected_url = endpoint + suffix + elif transport.endswith("_alias"): + model = "azure_ai/responses/gpt-4o" if transport == "responses_alias" else "azure_ai/gpt-4o" + if transport == "responses_alias": + state.response_kind = "responses" + state.expected_url = f"{endpoint}/openai/v1/responses?api-version=preview" + monkeypatch.setenv("AZURE_AI_API_BASE", endpoint) + monkeypatch.setenv("AZURE_API_VERSION", version) + if key: + monkeypatch.setenv("AZURE_AI_API_KEY", key) + settings = _make_settings() + state.expected_api_key = key if transport not in ("v1", "raw") else None + factory = MagicMock(return_value=lambda: "companion-token") + monkeypatch.setattr(azure.identity, "get_bearer_token_provider", factory) + # Compare native auth selection, not a reimplementation of its precedence. + monkeypatch.setattr(litellm, "api_key", None) + await litellm.acompletion( + model=model, messages=[{"role": "user", "content": "usr"}], api_base=endpoint, api_key=key, + **({} if transport.startswith("raw") or transport == "responses_alias" else {"api_version": version}), + ) + if key is None: + assert state.sent[-1].headers["Authorization"] == "Bearer companion-token" + factory.assert_called() + cache = LLMClientCache() + native_companion_auth.append(cache) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", cache) + monkeypatch.setattr(litellm, "api_key", "foreign-global-key" if guard_key else None) + handler = LiteLLMAIHandler() + await state.invoke(handler, entrypoint, model=model, guard_key=guard_key) + assert len(state.sent) == 2 + for name in ("Authorization", "api-key"): + assert state.sent[1].headers.get_list(name) == state.sent[0].headers.get_list(name) + + +@pytest.mark.parametrize("transport", ("sdk", "responses", "raw")) +@pytest.mark.parametrize("failure", (False, True)) +@pytest.mark.asyncio +async def test_tokenless_azure_nested_contexts(monkeypatch, native_azure_oidc, transport, failure): + monkeypatch.delenv("AZURE_AD_TOKEN") + handlers = [] + for secret in ("first-secret", "second-secret"): + monkeypatch.setenv("AZURE_CLIENT_SECRET", secret) + handlers.append((LiteLLMAIHandler(), secret)) + monkeypatch.delenv("AZURE_CLIENT_SECRET") + inner = LiteLLMAIHandler() + model = {"sdk": "azure/gpt-4o", "responses": "azure/responses/gpt-4o", "raw": "azure_ai/test-model"}[transport] + inner_model = "azure/responses/gpt-4o" if transport == "sdk" else "azure/gpt-4o" + context_var = (litellm_handler._azure_oidc_request if transport == "sdk" + else litellm_handler._azure_ad_responses_request) + barrier = asyncio.Event() + arrived = 0 + + async def run(handler, secret): + async def nested(**kwargs): + assert litellm_handler._azure_oidc_request.get() is None + assert litellm_handler._azure_ad_responses_request.get() is None + + async def completion(**kwargs): + nonlocal arrived + context = context_var.get() + assert context["dispatch_token"] is None + assert context["auth_environment"]["AZURE_CLIENT_SECRET"] == secret + arrived += 1 + if arrived == 2: + barrier.set() + await asyncio.wait_for(barrier.wait(), timeout=5) + await inner._acompletion(_completion=nested, model=inner_model) + assert context_var.get() is context + if failure: + raise RuntimeError("synthetic companion failure") + + try: + await handler._acompletion( + _completion=completion, model=model, api_key=DUMMY_LITELLM_API_KEY, + _azure_ad_guard=True, _raw_api_key_guard=True, ) + except RuntimeError as error: + assert failure and str(error) == "synthetic companion failure" + else: + assert not failure + assert litellm_handler._azure_oidc_request.get() is None + assert litellm_handler._azure_ad_responses_request.get() is None + + await asyncio.gather(*(run(handler, secret) for handler, secret in handlers)) + + +@pytest.mark.parametrize("host", ("https://owned.example", "https://owned.services.ai.azure.com")) +@pytest.mark.parametrize("headers", ({}, {"api-key": "gateway-key"}, {"Authorization": "Basic explicit"})) +@pytest.mark.asyncio +async def test_native_raw_azure_companion_snapshot(monkeypatch, native_azure_oidc, native_companion_auth, host, headers): + import azure.identity + + state = native_azure_oidc + monkeypatch.delenv("AZURE_AD_TOKEN") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "owned-secret") + monkeypatch.setenv("AZURE_AI_API_BASE", host) + settings = _make_settings({"LITELLM.EXTRA_HEADERS": json.dumps(headers)}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + handler = LiteLLMAIHandler() + for name in ("AZURE_CLIENT_SECRET", "AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_AI_API_KEY"): + monkeypatch.setenv(name, "foreign-value") + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://foreign.example") + monkeypatch.setenv("AZURE_AD_TOKEN", "oidc/env/FOREIGN_ASSERTION") + constructor = MagicMock(wraps=azure.identity.ClientSecretCredential) + monkeypatch.setattr(azure.identity, "ClientSecretCredential", constructor) + monkeypatch.setattr(azure.identity, "get_bearer_token_provider", lambda *args: lambda: "owned-companion-token") + state.expected_api_key = headers.get("api-key") + hostname = urlparse(host).hostname or "" + suffix = "/models/chat/completions" if hostname.endswith(".services.ai.azure.com") else "/chat/completions" + state.expected_url = host + suffix + expected = None if "api-key" in headers else "Bearer owned-companion-token" + assert await state.invoke(handler, model="azure_ai/test-model") == expected + if "api-key" in headers: + constructor.assert_not_called() + else: + assert constructor.call_args.args == ("owned-tenant", "owned-client", "owned-secret") + assert constructor.call_args.kwargs["authority"] == "https://owned-authority.example" + assert state.exchanged == [] + assert handler._request_headers == headers + + +@pytest.mark.parametrize("source", ("azure_key", "AZURE_API_KEY", "AZURE_OPENAI_API_KEY", "AZURE_AD_TOKEN")) +@pytest.mark.asyncio +async def test_native_raw_azure_companion_preserves_initial_auth( + monkeypatch, native_azure_oidc, native_companion_auth, source, +): + import azure.identity + + state = native_azure_oidc + monkeypatch.delenv("AZURE_AD_TOKEN") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "owned-secret") + monkeypatch.setenv("AZURE_AI_API_BASE", "https://owned.example") + monkeypatch.setattr(litellm_handler, "get_settings", _make_settings) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + if source == "azure_key": + monkeypatch.setattr(litellm, source, "initial-key") + else: + monkeypatch.setenv(source, "initial-token" if source == "AZURE_AD_TOKEN" else "initial-key") + monkeypatch.setattr(azure.identity, "get_bearer_token_provider", lambda *args: lambda: "companion-token") + state.expected_url = "https://owned.example/chat/completions" + state.expected_api_key = None if source == "AZURE_AD_TOKEN" else "initial-key" + handler = LiteLLMAIHandler() + assert not handler._uses_captured_azure_companion_auth("azure_ai") + assert await state.invoke(handler, model="azure_ai/test-model", guard_key=False) == ( + "Bearer companion-token" if source == "AZURE_AD_TOKEN" else None + ) + + +@pytest.mark.parametrize("field", ("AZURE_CLIENT_SECRET", "AZURE_AUTHORITY_HOST", "AZURE_SCOPE")) +@pytest.mark.parametrize("token", (None, "owned-ad-token")) +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.parametrize("refresh", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_ad_sdk_cache_and_refresh(monkeypatch, native_azure_oidc, field, token, guard_key, refresh): + import azure.identity - assert "api_key" not in mock_call.call_args[1], ( - f"A keyless init must clear the previous request's provider key; " - f"kwargs had: {mock_call.call_args[1].get('api_key')!r}" + state = native_azure_oidc + if token is None: + monkeypatch.delenv("AZURE_AD_TOKEN") + else: + monkeypatch.setenv("AZURE_AD_TOKEN", token) + monkeypatch.setenv("AZURE_CLIENT_SECRET", "owned-secret") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", refresh) + settings = _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": "https://owned.openai.azure.com", "OPENAI.API_VERSION": "v1", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + state.expected_url = "https://owned.openai.azure.com/openai/v1/chat/completions" + identities, calls = [], [] + + def credential(tenant, client, secret, *, authority): + assert (tenant, client) == ("owned-tenant", "owned-client") + return SimpleNamespace(secret=secret, authority=authority) + + def provider(credential, scope): + index = len(identities) + identities.append((credential.secret, credential.authority, scope)) + calls.append(0) + + def token(): + calls[index] += 1 + return f"identity-{index}-token-{calls[index]}" + + return token + + monkeypatch.setattr(azure.identity, "ClientSecretCredential", credential) + monkeypatch.setattr(azure.identity, "get_bearer_token_provider", provider) + first = LiteLLMAIHandler() + assert (await state.invoke(first, guard_key=guard_key)).startswith("Bearer identity-0-token-") + first_identity = identities[0] + changed = "second-secret" if field == "AZURE_CLIENT_SECRET" else "https://second.example" + monkeypatch.setenv(field, changed) + second = LiteLLMAIHandler() + for variable in ("AZURE_CLIENT_SECRET", "AZURE_AUTHORITY_HOST", "AZURE_SCOPE"): + monkeypatch.setenv(variable, "foreign-value") + assert (await state.invoke(second, guard_key=guard_key)).startswith("Bearer identity-1-token-") + assert (await state.invoke(first, guard_key=guard_key)).startswith("Bearer identity-0-token-") + expected_identity = list(first_identity) + expected_identity[("AZURE_CLIENT_SECRET", "AZURE_AUTHORITY_HOST", "AZURE_SCOPE").index(field)] = changed + assert identities == [first_identity, tuple(expected_identity)] + before = list(calls) + await asyncio.gather(state.invoke(first, guard_key=guard_key), state.invoke(second, guard_key=guard_key)) + assert {request.headers["Authorization"].split("-token-")[0] for request in state.sent[-2:]} == { + "Bearer identity-0", "Bearer identity-1", + } + assert all(after > previous for after, previous in zip(calls, before, strict=True)) + clients = [value for value in litellm.in_memory_llm_clients_cache.cache_dict.values() + if isinstance(value, openai.AsyncOpenAI)] + assert len(clients) == 2 + assert state.exchanged == [] + + +@pytest.mark.asyncio +async def test_ordinary_azure_sdk_context_does_not_select_oidc(monkeypatch, native_azure_oidc): + from litellm.llms.azure import common_utils as azure_common + + litellm_handler._install_azure_oidc_bridge() + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + monkeypatch.setattr(azure_common, "get_secret_str", lambda name: "outside-token" if name == "AZURE_AD_TOKEN" else None) + context = litellm_handler._azure_oidc_request.set({ + "selector": None, "dispatch_token": "owned-token", "generated_guard": True, + }) + try: + assert azure_common.get_azure_ad_token({}) == "outside-token" + params = azure_common.GenericLiteLLMParams(api_key=DUMMY_LITELLM_API_KEY, azure_ad_token="owned-token") + assert azure_common.BaseAzureLLM._base_validate_azure_environment({}, params) == { + "api-key": DUMMY_LITELLM_API_KEY, + } + finally: + litellm_handler._azure_oidc_request.reset(context) + + +@pytest.mark.parametrize("concurrent", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_ad_responses_context(monkeypatch, native_azure_oidc, concurrent): + from litellm.llms.azure import common_utils as azure_common + + state = native_azure_oidc + state.response_kind = "responses" + state.expected_url = "https://owned.openai.azure.com/openai/responses?api-version=2024-06-01" + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + handlers = [] + for token in ("first-token", "second-token"): + monkeypatch.setenv("AZURE_AD_TOKEN", token) + handlers.append(LiteLLMAIHandler()) + monkeypatch.setattr(litellm, "api_key", "foreign-key") + barrier = asyncio.Event() + arrived = 0 + native_send = httpx.AsyncClient.send + + async def send(client, request, **kwargs): + nonlocal arrived + arrived += 1 + if arrived == 2: + barrier.set() + if concurrent: + await barrier.wait() + return await native_send(client, request, **kwargs) + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + if concurrent: + await asyncio.gather(*(handler.probe_completion("azure/responses/gpt-4o") for handler in handlers)) + else: + for handler in handlers: + await handler.probe_completion("azure/responses/gpt-4o") + assert sorted(request.headers["authorization"] for request in state.sent) == [ + "Bearer first-token", "Bearer second-token", + ] + assert litellm_handler._azure_ad_responses_request.get() is None + params = azure_common.GenericLiteLLMParams(azure_ad_token="outside-token") + assert azure_common.BaseAzureLLM._base_validate_azure_environment({}, params) == {"api-key": "foreign-key"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("guard_key", (False, True)) +async def test_native_azure_ad_responses_rejects_implicit_discovery(monkeypatch, native_azure_oidc, guard_key): + from litellm.llms.azure import common_utils as azure_common + + state = native_azure_oidc + monkeypatch.setenv("AZURE_AD_TOKEN", "owned-ad-token") + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + provider = MagicMock(side_effect=AssertionError("Implicit discovery must not run")) + monkeypatch.setattr(azure_common, "get_azure_ad_token_provider", provider) + with pytest.raises(litellm.APIConnectionError, match="implicit Azure AD credential discovery"): + await state.invoke(handler, model="azure/responses/gpt-4o", guard_key=guard_key) + provider.assert_not_called() + assert state.sent == [] + assert litellm_handler._azure_ad_responses_request.get() is None + + +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_ad_responses_propagates_companion_failure(monkeypatch, native_azure_oidc, guard_key): + import azure.identity + + state = native_azure_oidc + monkeypatch.setenv("AZURE_AD_TOKEN", "owned-ad-token") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "owned-secret") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + handler = LiteLLMAIHandler() + provider = MagicMock(side_effect=RuntimeError("synthetic captured provider failure")) + monkeypatch.setattr(azure.identity, "get_bearer_token_provider", lambda *args: provider) + with pytest.raises(litellm.APIConnectionError, match="synthetic captured provider failure"): + await state.invoke(handler, model="azure/responses/gpt-4o", guard_key=guard_key) + provider.assert_called_once() + assert state.sent == state.exchanged == [] + assert litellm_handler._azure_ad_responses_request.get() is None + + +@pytest.mark.parametrize("transport", ("chat", "v1", "responses")) +@pytest.mark.parametrize("source", ("client_secret", "password")) +@pytest.mark.parametrize("timing", ("late", "initial")) +@pytest.mark.parametrize("guard_key", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_oidc_auth_sources(monkeypatch, native_azure_oidc, transport, source, timing, guard_key): + import azure.identity + + state = native_azure_oidc + if transport == "v1": + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": "https://owned.openai.azure.com", + "OPENAI.API_VERSION": "v1", + })) + state.expected_url = "https://owned.openai.azure.com/openai/v1/chat/completions" + if timing == "initial": + if source == "client_secret": + monkeypatch.setenv("AZURE_CLIENT_SECRET", "owned-secret") + else: + monkeypatch.setenv("AZURE_USERNAME", "owned-user") + monkeypatch.setenv("AZURE_PASSWORD", "owned-password") + handler = LiteLLMAIHandler() + if source == "client_secret": + monkeypatch.setenv("AZURE_CLIENT_SECRET", "foreign-secret") + credential_name = "ClientSecretCredential" + else: + monkeypatch.setenv("AZURE_USERNAME", "foreign-user") + monkeypatch.setenv("AZURE_PASSWORD", "foreign-password") + credential_name = "UsernamePasswordCredential" + constructor = MagicMock(wraps=getattr(azure.identity, credential_name)) + monkeypatch.setattr(azure.identity, credential_name, constructor) + provider = MagicMock(return_value=lambda: "provider-token") + monkeypatch.setattr(azure.identity, "get_bearer_token_provider", provider) + model = "azure/gpt-4o" + if transport == "responses": + state.response_kind = "responses" + state.expected_url = "https://owned.openai.azure.com/openai/responses?api-version=2024-06-01" + model = "azure/responses/gpt-4o" + expected = "provider-token" if timing == "initial" and transport in {"v1", "responses"} else "exchanged-token-1" + assert await state.invoke(handler, guard_key=guard_key, model=model) == f"Bearer {expected}" + if timing == "late": + provider.assert_not_called() + constructor.assert_not_called() + else: + provider.assert_called_once() + constructor.assert_called_once() + if source == "client_secret": + assert constructor.call_args.args == ("owned-tenant", "owned-client", "owned-secret") + else: + assert constructor.call_args.kwargs["username"] == "owned-user" + assert constructor.call_args.kwargs["password"] == "owned-password" + + +@pytest.mark.parametrize("transport", ("v1", "responses")) +@pytest.mark.parametrize("source", ("client_secret", "password")) +@pytest.mark.parametrize("warm", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_oidc_companion_authority(monkeypatch, native_azure_oidc, transport, source, warm): + from azure.core.credentials import AccessTokenInfo + from azure.identity import ClientSecretCredential, UsernamePasswordCredential + from litellm.llms.azure import common_utils as azure_common + + state = native_azure_oidc + client_id = f"owned-client-{transport}-{source}-{warm}" + monkeypatch.setenv("AZURE_CLIENT_ID", client_id) + if source == "client_secret": + monkeypatch.setenv("AZURE_CLIENT_SECRET", "owned-secret") + else: + monkeypatch.setenv("AZURE_USERNAME", "owned-user") + monkeypatch.setenv("AZURE_PASSWORD", "owned-password") + selected_authorities = [] + + def acquire_token(credential, *scopes, **kwargs): + # Keep native factories, credential construction and bearer-token policy; + # intercept acquisition only after Azure Identity has selected authority. + selected_authorities.append(credential._authority) + # Force policy refresh so a reused provider must retain its authority. + return AccessTokenInfo(f"provider-{credential._authority.split('//')[-1]}", 0) + + monkeypatch.setattr(ClientSecretCredential, "get_token_info", acquire_token) + monkeypatch.setattr(UsernamePasswordCredential, "get_token_info", acquire_token) + if warm and source == "client_secret": + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://foreign-authority.example") + # Seed the real native provider cache outside the request bridge. + outside_provider = azure_common.get_azure_ad_token_from_entra_id( + "owned-tenant", client_id, "owned-secret", state.inputs["AZURE_SCOPE"], ) + monkeypatch.setenv("AZURE_AUTHORITY_HOST", state.inputs["AZURE_AUTHORITY_HOST"]) + model = "azure/gpt-4o" + if transport == "v1": + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": "https://owned.openai.azure.com", + "OPENAI.API_VERSION": "v1", + })) + state.expected_url = "https://owned.openai.azure.com/openai/v1/chat/completions" + else: + state.response_kind = "responses" + state.expected_url = "https://owned.openai.azure.com/openai/responses?api-version=2024-06-01" + model = "azure/responses/gpt-4o" + first = LiteLLMAIHandler() + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://foreign-authority.example") + assert await state.invoke(first, model=model) == "Bearer provider-owned-authority.example" + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://second-authority.example") + second = LiteLLMAIHandler() + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://foreign-authority.example") + assert await state.invoke(second, model=model) == "Bearer provider-second-authority.example" + assert await state.invoke(first, model=model) == "Bearer provider-owned-authority.example" + assert set(selected_authorities) == {"https://owned-authority.example", "https://second-authority.example"} + if warm and source == "client_secret": + assert azure_common.get_azure_ad_token_from_entra_id( + "owned-tenant", client_id, "owned-secret", state.inputs["AZURE_SCOPE"], + ) is outside_provider - def test_placeholder_not_stored_in_global_litellm_api_key(self): - """With no OpenAI key configured, litellm.api_key must stay falsy.""" - LiteLLMAIHandler() - assert litellm.api_key != DUMMY_LITELLM_API_KEY, ( - "The placeholder must not occupy litellm.api_key — it shadows provider " - "env vars such as OPENROUTER_API_KEY in LiteLLM's resolution chain." - ) - assert litellm.openai_key == DUMMY_LITELLM_API_KEY, ( - "The placeholder must still be available on the OpenAI-compatible path, " - "so keyless local endpoints (vLLM, LM Studio, ...) keep working." - ) - - def test_placeholder_does_not_stick_across_handler_inits(self, monkeypatch): - """A keyless request must not poison later requests that do configure OPENAI.KEY. - - LiteLLMAIHandler is constructed per request but writes to LiteLLM's process-wide - globals. With the placeholder on litellm.api_key a single keyless init left - "dummy_key" there for the lifetime of the process, and every later request - resolved the placeholder even with OPENAI.KEY set — because __init__ only ever - writes the real OpenAI key to litellm.openai_key, which sits *behind* - litellm.api_key in the chain. Keeping the placeholder on openai_key means the - real key simply replaces it. - """ - from litellm import main as litellm_main - - openai_settings = type("Settings", (), { - "config": type("Config", (), { - "reasoning_effort": None, - "ai_timeout": 30, - "custom_reasoning_model": False, - "max_model_tokens": 32000, - "verbosity_level": 0, - "seed": -1, - "get": lambda self, key, default=None: default, - })(), - "litellm": type("LiteLLM", (), { - "get": lambda self, key, default=None: default, - })(), - "openai": type("OpenAI", (), {"key": "sk-real-openai"})(), - "get": lambda self, key, default=None: ( - "sk-real-openai" if key == "OPENAI.KEY" else default - ), - })() +@pytest.mark.parametrize("source", ("client_secret", "password")) +@pytest.mark.parametrize("captured", (None, "")) +@pytest.mark.asyncio +async def test_native_azure_oidc_companion_authority_default(monkeypatch, native_azure_oidc, source, captured): + from azure.core.credentials import AccessTokenInfo + from azure.identity import ClientSecretCredential, UsernamePasswordCredential + + state = native_azure_oidc + if captured is None: + monkeypatch.delenv("AZURE_AUTHORITY_HOST") + else: + monkeypatch.setenv("AZURE_AUTHORITY_HOST", captured) + if source == "client_secret": + monkeypatch.setenv("AZURE_CLIENT_SECRET", "owned-secret") + else: + monkeypatch.setenv("AZURE_USERNAME", "owned-user") + monkeypatch.setenv("AZURE_PASSWORD", "owned-password") + authorities = [] + + def acquire_token(credential, *scopes, **kwargs): + authorities.append(credential._authority) + return AccessTokenInfo("public-cloud-provider", 4102444800) - LiteLLMAIHandler() # request 1: no OpenAI key configured - monkeypatch.setattr(litellm_handler, "get_settings", lambda: openai_settings) - LiteLLMAIHandler() # request 2: OPENAI.KEY configured + monkeypatch.setattr(ClientSecretCredential, "get_token_info", acquire_token) + monkeypatch.setattr(UsernamePasswordCredential, "get_token_info", acquire_token) + state.response_kind = "responses" + state.expected_url = "https://owned.openai.azure.com/openai/responses?api-version=2024-06-01" + handler = LiteLLMAIHandler() + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://foreign-authority.example") + if captured == "": + with pytest.raises(litellm.APIConnectionError, match="nonempty captured authority"): + await state.invoke(handler, model="azure/responses/gpt-4o") + assert authorities == state.sent == state.exchanged == [] + else: + assert await state.invoke(handler, model="azure/responses/gpt-4o") == "Bearer public-cloud-provider" + assert authorities == ["https://login.microsoftonline.com"] - captured = {} - class _Capturing: - def completion(self, *args, **kwargs): - captured["api_key"] = kwargs.get("api_key") - raise _StopCall +@pytest.mark.parametrize("identity", ("same", "authority", "scope")) +@pytest.mark.parametrize("version", ("2024-06-01", "v1")) +@pytest.mark.asyncio +async def test_native_azure_oidc_bridge_sdk_cache(monkeypatch, native_azure_oidc, identity, version): + from litellm.llms.azure import common_utils as azure_common - monkeypatch.setattr(litellm_main, "openai_chat_completions", _Capturing()) - monkeypatch.setattr(litellm_main, "base_llm_http_handler", _Capturing()) - with contextlib.suppress(Exception): - litellm.completion(model="gpt-4o", messages=[{"role": "user", "content": "hi"}]) + state = native_azure_oidc + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + provider = MagicMock(side_effect=AssertionError("OIDC must not create an additional credential provider")) + monkeypatch.setattr(azure_common, "get_azure_ad_token_provider", provider) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({ + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": "https://owned.openai.azure.com", + "OPENAI.API_VERSION": version, + })) + if version == "v1": + state.expected_url = "https://owned.openai.azure.com/openai/v1/chat/completions" + first = LiteLLMAIHandler() + assert await state.invoke(first, guard_key=False) == "Bearer exchanged-token-1" + client_types = (openai.AsyncAzureOpenAI, openai.AsyncOpenAI) + clients = [value for value in litellm.in_memory_llm_clients_cache.cache_dict.values() + if isinstance(value, client_types)] + assert len(clients) == 1 + if identity == "same": + monkeypatch.delenv("PR_AGENT_TEST_OIDC_ASSERTION") + else: + field = "AZURE_AUTHORITY_HOST" if identity == "authority" else "AZURE_SCOPE" + monkeypatch.setenv(field, "https://second.example") + second = LiteLLMAIHandler() + expected_count = 1 if identity == "same" else 2 + assert await state.invoke(second, guard_key=False) == f"Bearer exchanged-token-{expected_count}" + assert await state.invoke(first, guard_key=False) == "Bearer exchanged-token-1" + cached = [value for value in litellm.in_memory_llm_clients_cache.cache_dict.values() + if isinstance(value, client_types)] + assert len(cached) == expected_count + assert any(value is clients[0] for value in cached) + assert len(state.exchanged) == expected_count + if identity == "authority": + assert state.exchanged[1][0] == "https://second.example/owned-tenant/oauth2/v2.0/token" + elif identity == "scope": + assert state.exchanged[1][1]["scope"] == "https://second.example" + provider.assert_not_called() - assert captured.get("api_key") == "sk-real-openai", ( - f"A keyless init must not leave a placeholder that outlives it; " - f"LiteLLM resolved: {captured.get('api_key')!r}" + +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.asyncio +async def test_native_azure_oidc_bridge_text(monkeypatch, native_azure_oidc, stream): + state = native_azure_oidc + state.response_kind, state.stream = "text", stream + model = "azure_text/gpt-3.5-turbo-instruct" + state.expected_url = ( + "https://owned.openai.azure.com/openai/deployments/gpt-3.5-turbo-instruct/completions?api-version=2024-06-01" + ) + handler = LiteLLMAIHandler() + monkeypatch.setenv("AZURE_TENANT_ID", "foreign-tenant") + params = handler._get_provider_request_params(model) + try: + response = await handler._acompletion( + model=params.pop("model", model), messages=[{"role": "user", "content": "test"}], stream=stream, **params, ) + finally: + assert litellm_handler._azure_oidc_request.get() is None + assert len(state.exchanged) == 1 + assert state.exchanged[0][0] == "https://owned-authority.example/owned-tenant/oauth2/v2.0/token" + assert state.exchanged[0][1]["client_id"] == "owned-client" + if stream: + monkeypatch.setenv("AZURE_CLIENT_ID", "foreign-after-return") + assert [chunk async for chunk in response] + assert len(state.sent) == len(state.exchanged) == 1 + assert state.sent[0].headers.get_list("authorization") == ["Bearer exchanged-token-1"] + assert litellm_handler._azure_oidc_request.get() is None - def test_openrouter_env_key_wins_over_placeholder(self, monkeypatch): - """End-to-end through LiteLLM's own resolution chain. - Patches LiteLLM's internal HTTP handler to capture the api_key that would be - sent for an ``openrouter/*`` model when the key comes from the native - OPENROUTER_API_KEY env var rather than PR-Agent's [openrouter] settings. - """ - from litellm import main as litellm_main +@pytest.mark.parametrize(("transport", "key_source"), ( + ("cloudflare", "setting"), ("cloudflare", "header"), ("responses", "setting"), +)) +@pytest.mark.parametrize("api_key", ("owned-key", "dummy_key")) +@pytest.mark.asyncio +async def test_native_azure_oidc_bridge_api_key_winner(monkeypatch, native_azure_oidc, transport, key_source, api_key): + state = native_azure_oidc + endpoint = "https://owned.openai.azure.com" + model = "azure/responses/gpt-4o" + if transport == "cloudflare": + endpoint = "https://gateway.ai.cloudflare.com/v1/account/gateway/azure-openai/resource" + model = "azure/gpt-4o" + state.expected_url = f"{endpoint}/gpt-4o/chat/completions?api-version=2024-06-01" + else: + state.response_kind = "responses" + state.expected_url = f"{endpoint}/openai/responses?api-version=2024-06-01" + settings = { + "OPENAI.API_TYPE": "azure", "OPENAI.API_BASE": endpoint, "OPENAI.API_VERSION": "2024-06-01", + "OPENAI.KEY": api_key if key_source == "setting" else None, + "LITELLM.EXTRA_HEADERS": json.dumps({"api-key": api_key}) if key_source == "header" else None, + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(settings)) + handler = LiteLLMAIHandler() + monkeypatch.delenv("PR_AGENT_TEST_OIDC_ASSERTION") + state.expected_api_key = api_key + assert await state.invoke(handler, guard_key=key_source == "header" and transport == "cloudflare", + model=model) is None + assert len(state.sent) == 1 + assert state.sent[0].headers.get_list("authorization") == [] + assert state.exchanged == [] - monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-real-key") - LiteLLMAIHandler() - captured = {} +@pytest.mark.parametrize("action", ("outside", "alias", "error")) +@pytest.mark.asyncio +async def test_native_azure_oidc_bridge_context(monkeypatch, native_azure_oidc, action): + from litellm.llms.azure import azure as azure_module + from litellm.llms.azure import common_utils as azure_common - class _CapturingHandler: - def completion(self, **kwargs): - captured["api_key"] = kwargs.get("api_key") - raise _StopCall + state = native_azure_oidc + handler = LiteLLMAIHandler() + assert await state.invoke(handler, guard_key=False) == "Bearer exchanged-token-1" + assert litellm_handler._azure_oidc_request.get() is None + if action == "outside": + monkeypatch.setenv("AZURE_TENANT_ID", "outside-tenant") + state.ambient = dict(os.environ) + token = await asyncio.to_thread( + azure_common.get_azure_ad_token_from_oidc, "oidc/env/PR_AGENT_TEST_OIDC_ASSERTION", + "outside-client", None, "outside-scope", + ) + assert token == "exchanged-token-2" + assert state.exchanged[1][0] == "https://owned-authority.example/outside-tenant/oauth2/v2.0/token" + assert state.exchanged[1][1]["client_id"] == "outside-client" + assert state.exchanged[1][1]["scope"] == "outside-scope" + elif action == "alias": + monkeypatch.setattr(azure_module, "get_azure_ad_token_from_oidc", lambda *args, **kwargs: "unsafe") + with pytest.raises(RuntimeError, match="(?i)replaced|incompatible"): + await state.invoke(handler, guard_key=False) + assert len(state.exchanged) == 1 + else: + monkeypatch.setenv("AZURE_SCOPE", "second-scope") + second = LiteLLMAIHandler() + state.status = 401 + with pytest.raises(litellm.AuthenticationError): + await state.invoke(second, guard_key=False) + assert len(state.exchanged) == 2 + assert len(state.sent) == 1 + assert litellm_handler._azure_oidc_request.get() is None - monkeypatch.setattr(litellm_main, "base_llm_http_handler", _CapturingHandler()) - # LiteLLM maps whatever the transport raises onto its own exception types, - # so the sentinel comes back wrapped — the captured key is what matters. - with contextlib.suppress(Exception): - litellm.completion( - model="openrouter/deepseek/deepseek-chat", - messages=[{"role": "user", "content": "hi"}], - ) - assert captured.get("api_key") == "sk-or-real-key", ( - f"OPENROUTER_API_KEY must not be shadowed by the placeholder; " - f"LiteLLM resolved: {captured.get('api_key')!r}" +async def _call(handler, model): + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + await handler.chat_completion(model=model, system="sys", user="usr") + return mock_call.call_args.kwargs + + +@pytest.mark.parametrize(("module_name", "symbol"), ( + ("litellm.llms.anthropic.common_utils", "AnthropicModelInfo"), + ("litellm.llms.openai_like.json_loader", "JSONProviderRegistry"), + ("litellm.utils", "_get_model_info_helper"), +)) +@pytest.mark.asyncio +async def test_missing_private_import_preserves_module_and_fails_closed(monkeypatch, module_name, symbol): + import builtins + import importlib.util + + spec = importlib.util.spec_from_file_location("_isolated_litellm_handler", litellm_handler.__file__) + module = importlib.util.module_from_spec(spec) + original_import = builtins.__import__ + + def guarded_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == module_name and symbol in fromlist and (globals or {}).get("__name__") == spec.name: + raise ImportError("Unavailable private interface") + return original_import(name, globals, locals, fromlist, level) + + with monkeypatch.context() as scoped: + scoped.setattr(builtins, "__import__", guarded_import) + spec.loader.exec_module(module) + assert getattr(module, symbol) is None + monkeypatch.setattr(module, "get_settings", lambda: _make_settings()) + if symbol == "JSONProviderRegistry": + with pytest.raises(RuntimeError, match="JSONProviderRegistry.*request isolation"): + module.LiteLLMAIHandler() + return + handler = module.LiteLLMAIHandler() + completion = AsyncMock(return_value=_mock_response()) + monkeypatch.setattr(module, "acompletion", completion) + if symbol == "AnthropicModelInfo": + with pytest.raises(RuntimeError, match="AnthropicModelInfo.*request isolation"): + await handler.probe_completion("anthropic/claude-sonnet-4") + completion.assert_not_called() + await handler.probe_completion("gpt-4o") + else: + with pytest.raises(RuntimeError, match="_get_model_info_helper.*request isolation"): + await handler.probe_completion("gpt-4o") + completion.assert_not_called() + assert module._uses_openai_responses_transport("openai/responses/gpt-4o", "openai") + assert not module._uses_openai_responses_transport("ft:babbage-002:example", "openai") + await handler.probe_completion("anthropic/claude-sonnet-4") + + +@pytest.mark.parametrize("method", ("list_providers", "get", "exists")) +def test_incomplete_registry_fails_before_snapshot(monkeypatch, method): + registry = type("Registry", (), { + name: staticmethod(lambda *args: ()) for name in ("list_providers", "get", "exists") + }) + setattr(registry, method, None) + with monkeypatch.context() as scoped: + scoped.setattr(litellm_handler, "JSONProviderRegistry", registry) + settings = MagicMock(side_effect=AssertionError("Do not start constructing snapshots")) + scoped.setattr(litellm_handler, "get_settings", settings) + with pytest.raises(RuntimeError, match="JSONProviderRegistry"): + LiteLLMAIHandler() + settings.assert_not_called() + + +@pytest.mark.parametrize("method", ("get_api_key", "get_auth_token")) +def test_incomplete_anthropic_interface_does_not_install_partial_bridge(monkeypatch, method): + info = type("ModelInfo", (), { + name: staticmethod(lambda value=None: value) for name in ("get_api_key", "get_auth_token") + }) + setattr(info, method, None) + originals = (info.get_api_key, info.get_auth_token) + with monkeypatch.context() as scoped: + scoped.setattr(litellm_handler, "AnthropicModelInfo", info) + with pytest.raises(RuntimeError, match="AnthropicModelInfo"): + litellm_handler._install_anthropic_auth_token_bridge() + assert (info.get_api_key, info.get_auth_token) == originals + + +def test_noncallable_model_info_helper_does_not_silently_choose_transport(monkeypatch): + monkeypatch.setattr(litellm_handler, "_get_model_info_helper", object()) + with pytest.raises(RuntimeError, match="_get_model_info_helper"): + litellm_handler._uses_openai_responses_transport("gpt-4o", "openai") + + +def test_provider_environment_tables_cover_known_or_legacy_transports(): + # LiteLLM retains explicit transports for these providers outside provider_list. + legacy_transports = {"aleph_alpha", "anyscale"} + known_providers = set(litellm.provider_list) | set(litellm_handler.JSONProviderRegistry.list_providers()) + for table in (litellm_handler.PROVIDER_API_KEY_ENV_VARS, litellm_handler.PROVIDER_API_BASE_ENV_VARS): + assert set(table) <= known_providers | legacy_transports + + +@pytest.mark.parametrize("provider", ("aleph_alpha", "anyscale")) +@pytest.mark.parametrize("initial_key", (None, "handler-key")) +@pytest.mark.asyncio +async def test_legacy_provider_transport_uses_request_local_credentials(monkeypatch, provider, initial_key): + from litellm.litellm_core_utils import logging_worker + + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + monkeypatch.setattr(openai, "organization", None) + settings = _make_settings() + settings.litellm.custom_llm_provider = provider + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + key_variable = "ALEPH_ALPHA_API_KEY" if provider == "aleph_alpha" else "ANYSCALE_API_KEY" + base_variable = "ALEPH_ALPHA_API_BASE" if provider == "aleph_alpha" else "ANYSCALE_API_BASE" + api_base = "https://handler.example/v1" + monkeypatch.setenv(base_variable, api_base) + if initial_key: + monkeypatch.setenv(key_variable, initial_key) + handler = LiteLLMAIHandler() + monkeypatch.setenv(key_variable, "another-handler-key") + monkeypatch.setenv(base_variable, "https://another-handler.example/v1") + monkeypatch.setattr(litellm, "api_key", "another-global-key") + monkeypatch.setattr(litellm, "aleph_alpha_key", "another-aleph-key") + captured = [] + + def respond(request): + captured.append(request) + if provider == "aleph_alpha": + payload = {"completions": [{"completion": "ok", "finish_reason": "stop"}]} + else: + payload = {"choices": [{"message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}]} + return httpx.Response(200, request=request, json=payload) + + def post(url, *, headers, data, stream): + return respond(httpx.Request("POST", url, headers=headers, content=data)) + + async def send(client, request, **kwargs): + return respond(request) + + # Keep PR-Agent and LiteLLM dispatch/authentication real; intercept only HTTP. + monkeypatch.setattr(litellm.module_level_client, "post", post) + monkeypatch.setattr(httpx.AsyncClient, "send", send) + try: + content, finish_reason = await handler.chat_completion(model="model", system="sys", user="usr") + finally: + try: + # LiteLLM schedules the task that enqueues logging after completion. + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + + assert (content, finish_reason) == ("ok", "stop") + assert len(captured) == 1 + expected_url = api_base if provider == "aleph_alpha" else f"{api_base}/chat/completions" + assert str(captured[0].url) == expected_url + assert captured[0].headers["authorization"] == f"Bearer {initial_key or DUMMY_LITELLM_API_KEY}" + + +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("extra_headers", (None, {"X-Request": "handler-value"})) +@pytest.mark.asyncio +async def test_together_transport_preserves_request_headers(monkeypatch, entrypoint, extra_headers): + from litellm.litellm_core_utils import logging_worker + + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + settings = _make_settings({"LITELLM.EXTRA_HEADERS": json.dumps(extra_headers) if extra_headers else None}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + monkeypatch.setenv("TOGETHERAI_API_KEY", "handler-key") + monkeypatch.setenv("TOGETHER_AI_API_BASE", "https://together-handler.example/v1") + handler = LiteLLMAIHandler() + monkeypatch.setenv("TOGETHERAI_API_KEY", "another-handler-key") + monkeypatch.setenv("TOGETHER_AI_API_BASE", "https://another-handler.example/v1") + captured = [] + + async def send(client, request, **kwargs): + captured.append(request) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "test", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + + # Exercise native header assembly: SDK Omit values must never reach httpx. + monkeypatch.setattr(httpx.AsyncClient, "send", send) + model = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo" + try: + if entrypoint == "chat": + assert await handler.chat_completion(model, "sys", "usr") == ("ok", "stop") + else: + await handler.probe_completion(model) + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + + assert len(captured) == 1 + request = captured[0] + assert str(request.url) == "https://together-handler.example/v1/chat/completions" + assert request.headers["authorization"] == "Bearer handler-key" + assert "openai-organization" not in request.headers + assert "openai-project" not in request.headers + assert request.headers.get("x-request") == ("handler-value" if extra_headers else None) + + +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("extra_headers", (None, {"X-Request": "handler-value"}, { + "authorization": "Bearer explicit-handler-token", "editor-version": "request-editor", +})) +@pytest.mark.asyncio +async def test_copilot_transport_preserves_native_headers(monkeypatch, tmp_path, entrypoint, extra_headers): + from litellm.litellm_core_utils import logging_worker + from litellm.llms.github_copilot.authenticator import Authenticator + from litellm.llms.github_copilot.common_utils import EDITOR_PLUGIN_VERSION + + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path / "copilot")) + monkeypatch.setenv("GITHUB_COPILOT_API_BASE", "https://copilot-handler.example/v1") + monkeypatch.setattr(Authenticator, "get_api_key", lambda self: "handler-copilot-token") + monkeypatch.setattr(Authenticator, "get_api_base", lambda self: "https://copilot-handler.example/v1") + monkeypatch.delenv("OPENAI_CUSTOM_HEADERS", raising=False) + settings = _make_settings({"LITELLM.EXTRA_HEADERS": json.dumps(extra_headers) if extra_headers else None}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + handler = LiteLLMAIHandler() + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + captured = [] + + async def send(client, request, **kwargs): + captured.append(request) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + + # Keep native Copilot defaults and SDK assembly; stub token acquisition and HTTP only. + monkeypatch.setattr(httpx.AsyncClient, "send", send) + try: + if entrypoint == "chat": + assert await handler.chat_completion("github_copilot/gpt-4o", "sys", "usr") == ("ok", "stop") + else: + await handler.probe_completion("github_copilot/gpt-4o") + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + + assert len(captured) == 1 + request = captured[0] + assert str(request.url) == "https://copilot-handler.example/v1/chat/completions" + explicit = extra_headers or {} + assert request.headers.get_list("authorization") == [ + explicit.get("authorization", "Bearer handler-copilot-token") + ] + assert request.headers["copilot-integration-id"] == "vscode-chat" + assert request.headers["editor-version"] == explicit.get("editor-version", "vscode/1.95.0") + assert request.headers["editor-plugin-version"] == EDITOR_PLUGIN_VERSION + assert request.headers["openai-intent"] == "conversation-panel" + assert request.headers.get("x-request") == explicit.get("X-Request") + + +@pytest.mark.parametrize("provider", ("openai", "vercel_ai_gateway")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize( + "scenario", ( + "initial_accounts", "late_accounts", "initial_headers", "late_headers", "warm", "warm_accounts", "retry", + "custom_accounts", "custom_accounts_over_env", "explicit_headers", + ), +) +@pytest.mark.asyncio +async def test_native_sdk_header_snapshot(monkeypatch, provider, entrypoint, scenario): + from litellm.litellm_core_utils import logging_worker + from litellm.llms.openai.openai import OpenAIChatCompletion + + endpoint = f"https://{provider.replace('_', '-')}-{entrypoint}-{scenario.replace('_', '-')}.example/v1" + overrides = {"OPENAI.KEY": "handler-key", "OPENAI.API_BASE": endpoint} if provider == "openai" else {} + if provider == "vercel_ai_gateway": + monkeypatch.setenv("VERCEL_AI_GATEWAY_API_KEY", "handler-key") + monkeypatch.setenv("VERCEL_AI_GATEWAY_API_BASE", endpoint) + if scenario == "explicit_headers": + overrides["LITELLM.EXTRA_HEADERS"] = json.dumps({ + "authorization": "Bearer explicit-key", "openai-organization": "explicit-org", + "openai-project": "explicit-project", "X-Owned": "explicit-value", + }) + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer native-key\nX-Owned: native-value") + monkeypatch.setenv("OPENAI_ORG_ID", "native-org") + monkeypatch.setenv("OPENAI_PROJECT_ID", "native-project") + settings = _make_settings(overrides) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + if scenario in ("initial_accounts", "custom_accounts_over_env"): + monkeypatch.setenv("OPENAI_ORG_ID", "handler-org") + monkeypatch.setenv("OPENAI_PROJECT_ID", "handler-project") + if scenario == "initial_headers": + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "authorization: Bearer handler-header\nX-Owned: value") + if scenario in ("custom_accounts", "custom_accounts_over_env"): + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "openai-organization: custom-org\nOpenAI-Project: custom-project") + handler = LiteLLMAIHandler() + if scenario in ("late_headers", "warm", "retry"): + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer another-handler\nX-Foreign: value") + if scenario in ("late_accounts", "warm_accounts"): + monkeypatch.setenv("OPENAI_ORG_ID", "another-org") + monkeypatch.setenv("OPENAI_PROJECT_ID", "another-project") + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + captured, clients = [], [] + native_factory = OpenAIChatCompletion._get_openai_client + + def observe_client(self, *args, **kwargs): + client = native_factory(self, *args, **kwargs) + clients.append(client) + return client + + monkeypatch.setattr(OpenAIChatCompletion, "_get_openai_client", observe_client) + + async def send(client, request, **kwargs): + captured.append(request) + if scenario == "retry" and len(captured) == 1: + return httpx.Response(500, request=request, json={"error": {"message": "retry", "type": "server_error"}}) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + try: + for _ in range(2 if scenario in ("warm", "warm_accounts") else 1): + if entrypoint == "chat": + await handler.chat_completion(f"{provider}/gpt-4o", "sys", "usr") + else: + await handler.probe_completion(f"{provider}/gpt-4o") + if scenario == "warm": + monkeypatch.delenv("OPENAI_CUSTOM_HEADERS", raising=False) + if scenario == "warm_accounts": + monkeypatch.delenv("OPENAI_ORG_ID", raising=False) + monkeypatch.delenv("OPENAI_PROJECT_ID", raising=False) + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + + assert len(captured) == (2 if scenario in ("warm", "warm_accounts", "retry") else 1) + for request in captured: + assert str(request.url) == f"{endpoint}/chat/completions" + assert request.headers.get_list("authorization") == [ + "Bearer explicit-key" if scenario == "explicit_headers" + else "Bearer handler-header" if scenario == "initial_headers" else "Bearer handler-key" + ] + custom_accounts = scenario in ("custom_accounts", "custom_accounts_over_env") + assert request.headers.get("openai-organization") == ( + "explicit-org" if scenario == "explicit_headers" + else "custom-org" if custom_accounts else "handler-org" if scenario == "initial_accounts" else None + ) + assert request.headers.get("openai-project") == ( + "explicit-project" if scenario == "explicit_headers" + else "custom-project" if custom_accounts else "handler-project" if scenario == "initial_accounts" else None ) + assert request.headers.get("x-owned") == ( + "explicit-value" if scenario == "explicit_headers" else "value" if scenario == "initial_headers" else None + ) + assert "x-foreign" not in request.headers + assert litellm_handler._SDK_HEADER_MARKER not in request.headers + if scenario == "warm": + assert clients[0] is clients[1] + # The cached client's foreign defaults stay untouched; only request headers are isolated. + assert clients[0]._custom_headers["Authorization"] == "Bearer another-handler" + if scenario == "warm_accounts": + assert clients[0] is clients[1] + assert clients[0].organization == "another-org" + assert clients[0].project == "another-project" - def test_openai_compatible_endpoint_still_gets_placeholder(self, monkeypatch): - """Keyless local endpoints must still receive a key. - The OpenAI SDK raises "The api_key client option must be set" when no key is - resolved, which is exactly why the placeholder exists. Dropping it entirely - (instead of moving it to litellm.openai_key) would break these setups. - """ - from litellm import main as litellm_main +@pytest.mark.parametrize("azure", (False, True)) +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize("experimental", (False, True)) +@pytest.mark.asyncio +async def test_native_text_sdk_headers_survive_delayed_stream(monkeypatch, azure, stream, experimental): + from litellm.litellm_core_utils import logging_worker - LiteLLMAIHandler() + settings = _make_settings({ + "OPENAI.KEY": "handler-key", "OPENAI.API_BASE": "https://text-handler.example/v1", + "OPENAI.API_TYPE": "azure" if azure else None, "OPENAI.API_VERSION": "2024-02-15-preview", + "LITELLM.EXTRA_HEADERS": '{"X-Request": "handler-value"}', + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + if experimental: + monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") + monkeypatch.setenv("OPENAI_ORG_ID", "handler-org") + monkeypatch.setenv("OPENAI_PROJECT_ID", "handler-project") + handler = LiteLLMAIHandler() + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer foreign-token\nX-Foreign: value") + monkeypatch.setenv("OPENAI_ORG_ID", "foreign-org") + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + captured = [] - captured = {} - - class _CapturingOpenAI: - def completion(self, *args, **kwargs): - captured["api_key"] = kwargs.get("api_key") - raise _StopCall - - # LiteLLM routes the OpenAI path through either handler depending on the - # EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER flag; capture on both. - monkeypatch.setattr(litellm_main, "openai_chat_completions", _CapturingOpenAI()) - monkeypatch.setattr(litellm_main, "base_llm_http_handler", _CapturingOpenAI()) - with contextlib.suppress(Exception): - litellm.completion( - model="openai/local-model", - api_base="http://localhost:8000/v1", - messages=[{"role": "user", "content": "hi"}], - ) + async def send(client, request, **kwargs): + captured.append(request) + payload = { + "id": "test", "object": "text_completion", "created": 0, "model": "gpt-3.5-turbo-instruct", + "choices": [{"index": 0, "text": "ok", "logprobs": None, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + if stream: + # No stream_options.include_usage was requested for this text stream. + payload.pop("usage") + return httpx.Response(200, request=request, headers={"content-type": "text/event-stream"}, + content=f"data: {json.dumps(payload)}\n\ndata: [DONE]\n\n".encode()) + return httpx.Response(200, request=request, json=payload) - assert captured.get("api_key") == DUMMY_LITELLM_API_KEY, ( - f"Keyless OpenAI-compatible endpoints must still receive the placeholder; " - f"LiteLLM resolved: {captured.get('api_key')!r}" + monkeypatch.setattr(httpx.AsyncClient, "send", send) + model = ("azure_text/" if azure else "text-completion-openai/") + "gpt-3.5-turbo-instruct" + try: + params = handler._get_provider_request_params(model) + request_model = params.pop("model", model) + response = await handler._acompletion( + model=request_model, messages=[{"role": "user", "content": "test"}], stream=stream, **params, ) + assert litellm_handler._sdk_request_headers.get() is None + if stream: + monkeypatch.setenv("OPENAI_PROJECT_ID", "foreign-project-after-return") + chunks = [chunk async for chunk in response] + assert chunks + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + + assert len(captured) == 1 + request = captured[0] + assert request.url.host == "text-handler.example" + assert request.url.path.endswith("/completions") + assert request.headers.get("api-key" if azure else "authorization") == ( + "handler-key" if azure else "Bearer handler-key" + ) + assert request.headers["openai-organization"] == "handler-org" + assert request.headers["openai-project"] == "handler-project" + assert request.headers["x-request"] == "handler-value" + assert "x-foreign" not in request.headers + assert litellm_handler._SDK_HEADER_MARKER not in request.headers + + +@pytest.mark.asyncio +async def test_concurrent_sdk_requests_isolate_cached_default_headers(monkeypatch): + from litellm.litellm_core_utils import logging_worker + from litellm.llms.openai.openai import OpenAIChatCompletion + + settings = _make_settings({ + "OPENAI.KEY": "shared-handler-key", "OPENAI.API_BASE": "https://concurrent-sdk.example/v1", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + handlers = [] + for identity in ("one", "two"): + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", f"authorization: Bearer {identity}\nX-Owner: {identity}") + handlers.append(LiteLLMAIHandler()) + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer foreign\nX-Owner: foreign") + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + captured, clients = [], [] + entered = asyncio.Event() + native_factory = OpenAIChatCompletion._get_openai_client + + def observe_client(self, *args, **kwargs): + client = native_factory(self, *args, **kwargs) + clients.append(client) + return client + + async def send(client, request, **kwargs): + captured.append(request) + if len(captured) == 2: + entered.set() + await asyncio.wait_for(entered.wait(), 5) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + + monkeypatch.setattr(OpenAIChatCompletion, "_get_openai_client", observe_client) + monkeypatch.setattr(httpx.AsyncClient, "send", send) + try: + await asyncio.gather(*(handler.probe_completion("openai/gpt-4o") for handler in handlers)) + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + assert len(captured) == 2 + assert {tuple(request.headers.get_list("authorization")) for request in captured} == { + ("Bearer one",), ("Bearer two",), + } + for request in captured: + assert request.headers["authorization"] == f"Bearer {request.headers['x-owner']}" + assert litellm_handler._SDK_HEADER_MARKER not in request.headers + assert clients[0] is clients[1] + assert litellm_handler._sdk_request_headers.get() is None + + +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("api_version", ("2024-02-15-preview", "v1")) +@pytest.mark.parametrize("ad_auth", (False, True)) +@pytest.mark.asyncio +async def test_azure_deployment_selects_sdk_header_boundary(monkeypatch, entrypoint, api_version, ad_auth): + from litellm.litellm_core_utils import logging_worker + + settings = _make_settings({ + "OPENAI.KEY": None if ad_auth else "handler-key", "OPENAI.API_TYPE": "azure", + "OPENAI.API_BASE": "https://deployment-handler.example", "OPENAI.API_VERSION": api_version, + "OPENAI.DEPLOYMENT_ID": "custom-deployment", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + if ad_auth: + monkeypatch.setenv("AZURE_AD_TOKEN", "handler-ad-token") + handler = LiteLLMAIHandler() + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer foreign\nX-Foreign: value") + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + captured = [] + + async def send(client, request, **kwargs): + captured.append(request) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "custom-deployment", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + try: + if entrypoint == "chat": + await handler.chat_completion("azure/gpt-5-pro", "sys", "usr") + else: + await handler.probe_completion("azure/gpt-5-pro") + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + assert len(captured) == 1 + request = captured[0] + assert request.url.host == "deployment-handler.example" + assert request.url.path.endswith("/chat/completions") + if ad_auth or api_version == "v1": + assert request.headers.get_list("authorization") == [ + "Bearer handler-ad-token" if ad_auth else "Bearer handler-key" + ] + else: + assert request.headers["api-key"] == "handler-key" + assert "authorization" not in request.headers + assert "x-foreign" not in request.headers + assert litellm_handler._SDK_HEADER_MARKER not in request.headers + + +@pytest.mark.parametrize("custom_value", ("custom-account", "")) +@pytest.mark.parametrize("override", (None, "explicit-account", "omit")) +def test_sdk_account_defaults_preserve_explicit_override_and_omit(custom_value, override): + litellm_handler._install_sdk_header_bridge() + captured = [] + headers = {"OpenAI-Organization": custom_value, "OpenAI-Project": custom_value} + explicit = {} if override is None else { + name.lower(): openai.Omit() if override == "omit" else override for name in headers + } + snapshot = { + "organization": "env-org", "project": "env-project", "custom_headers": headers, + "explicit_headers": explicit, + } + request_headers = { + name: litellm_handler._CapturedSDKHeader() for name in headers + } + request_headers[litellm_handler._SDK_HEADER_MARKER] = litellm_handler._SDKHeaderSnapshot(snapshot) + + def send(request): + captured.append(request) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "gpt-4o", "choices": [], + }) + + with httpx.Client(transport=httpx.MockTransport(send)) as http_client: + with openai.OpenAI(api_key="handler-key", http_client=http_client) as client: + client.chat.completions.create(model="gpt-4o", messages=[], extra_headers=request_headers) + assert len(captured) == 1 + for name in headers: + expected = [] if override == "omit" else [custom_value if override is None else override] + assert captured[0].headers.get_list(name) == expected + assert litellm_handler._SDK_HEADER_MARKER not in captured[0].headers + + +def test_unmarked_sdk_request_keeps_native_headers_inside_handler_context(monkeypatch): + litellm_handler._install_sdk_header_bridge() + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer callback-token") + token = litellm_handler._sdk_request_headers.set({"custom_headers": {"Authorization": "Bearer handler-token"}}) + captured = [] + + def send(request): + captured.append(request) + return httpx.Response(200, request=request, json={ + "id": "callback", "object": "chat.completion", "created": 0, "model": "gpt-4o", "choices": [], + }) + + try: + with httpx.Client(transport=httpx.MockTransport(send)) as http_client: + with openai.OpenAI(api_key="callback-key", http_client=http_client) as client: + client.chat.completions.create(model="gpt-4o", messages=[]) + finally: + litellm_handler._sdk_request_headers.reset(token) + assert len(captured) == 1 + assert captured[0].headers.get_list("authorization") == ["Bearer callback-token"] + + +def test_sdk_header_logging_proxy_does_not_publish_or_mutate_marker(): + marker = litellm_handler._SDKHeaderSnapshot({"custom_headers": {"Authorization": "synthetic-secret"}}) + data = {"extra_headers": {litellm_handler._SDK_HEADER_MARKER: marker, "X-Request": "value"}} + logger = MagicMock() + proxy = litellm_handler._SDKHeaderLoggingProxy(logger) + proxy.pre_call(input="test", additional_args={"complete_input_dict": data}) + logged = logger.pre_call.call_args.kwargs["additional_args"]["complete_input_dict"] + assert logged == {"extra_headers": {"X-Request": "value"}} + assert data["extra_headers"][litellm_handler._SDK_HEADER_MARKER] is marker + proxy._deferred_stream_complete_args = ("response", False) + assert logger._deferred_stream_complete_args == ("response", False) + assert "synthetic-secret" not in repr(marker) + assert "synthetic-secret" not in str(marker) + + +@pytest.mark.asyncio +async def test_sdk_header_context_resets_after_cancellation(): + async def cancelled(**kwargs): + assert litellm_handler._sdk_request_headers.get() is not None + raise asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + await LiteLLMAIHandler()._acompletion(_completion=cancelled, model="openai/gpt-4o") + assert litellm_handler._sdk_request_headers.get() is None + + +def test_sdk_header_marker_fails_before_http_without_bridge(monkeypatch): + from openai._base_client import BaseClient + + litellm_handler._install_sdk_header_bridge() + native = BaseClient._build_headers._pr_agent_sdk_headers_original + monkeypatch.setattr(BaseClient, "_build_headers", native) + captured = [] + + def send(request): + captured.append(request) + raise AssertionError("An unconsumed marker must fail before HTTP") + + marker = litellm_handler._SDKHeaderSnapshot({"custom_headers": {"Authorization": "synthetic-secret"}}) + assert "synthetic-secret" not in repr(marker) + with httpx.Client(transport=httpx.MockTransport(send)) as http_client: + with openai.OpenAI(api_key="handler-key", http_client=http_client, max_retries=0) as client: + with pytest.raises(TypeError, match="Header value must be str or bytes"): + client.chat.completions.create(model="gpt-4o", messages=[], extra_headers={ + litellm_handler._SDK_HEADER_MARKER: marker, + }) + assert not captured + + +@pytest.mark.parametrize("initial_key", (None, "handler-key")) +@pytest.mark.asyncio +async def test_manus_responses_transport_uses_request_local_credentials(monkeypatch, initial_key): + from litellm.litellm_core_utils import logging_worker + + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + monkeypatch.delenv("MANUS_API_KEY", raising=False) + if initial_key: + monkeypatch.setenv("MANUS_API_KEY", initial_key) + monkeypatch.setenv("MANUS_API_BASE", "https://handler.example/v1") + handler = LiteLLMAIHandler() + monkeypatch.setenv("MANUS_API_KEY", "another-handler-key") + captured = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + captured.append(request) + raise TransportReached + + # Keep native completion-to-Responses dispatch and authentication intact. + monkeypatch.setattr(httpx.AsyncClient, "send", send) + model = "manus/responses/manus-1.6" + try: + with pytest.raises(TransportReached): + await handler._acompletion( + model=model, messages=[{"role": "user", "content": "test"}], + **handler._get_provider_request_params(model), + ) + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + assert len(captured) == 1 + assert str(captured[0].url) == "https://handler.example/v1/responses" + assert captured[0].headers["api_key"] == (initial_key or DUMMY_LITELLM_API_KEY) + + +@pytest.mark.parametrize("initial_key", (None, "handler-palm-key")) +@pytest.mark.asyncio +async def test_gemini_palm_transport_uses_request_local_credentials(monkeypatch, initial_key): + from litellm.litellm_core_utils import logging_worker + + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + for variable in ("GOOGLE_API_KEY", "GEMINI_API_KEY", "PALM_API_KEY"): + monkeypatch.delenv(variable, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + if initial_key: + monkeypatch.setenv("PALM_API_KEY", initial_key) + handler = LiteLLMAIHandler() + monkeypatch.setenv("PALM_API_KEY", "another-handler-palm-key") + captured = [] + + async def send(client, request, **kwargs): + captured.append(request) + return httpx.Response(200, request=request, json={ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "STOP", + "index": 0, + }], + "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3}, + }) + + # Keep native Gemini dispatch and authentication; intercept only HTTP. + monkeypatch.setattr(httpx.AsyncClient, "send", send) + try: + result = await handler.chat_completion(model="gemini/gemini-2.5-pro", system="sys", user="usr") + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + + assert result == ("ok", "stop") + assert len(captured) == 1 + assert captured[0].headers["x-goog-api-key"] == (initial_key or DUMMY_LITELLM_API_KEY) + + +@pytest.mark.parametrize("model", ("gemini-2.5-pro", "gemini-3-flash-preview")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe", "stream")) +@pytest.mark.parametrize("initial_key", (None, "owned-key")) +@pytest.mark.parametrize("initial_base", (None, "https://owned.example/v1")) +@pytest.mark.parametrize(("source", "header_key"), ( + ("environment", None), ("global", None), ("environment", "explicit-header-key"), +)) +@pytest.mark.asyncio +async def test_native_gemini_executor_boundary( + monkeypatch, model, entrypoint, initial_key, initial_base, source, header_key, +): + from litellm.litellm_core_utils import logging_worker + + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + if initial_key: + monkeypatch.setenv("GEMINI_API_KEY", initial_key) + if initial_base: + monkeypatch.setenv("GEMINI_API_BASE", initial_base) + if header_key: + settings = _make_settings({"LITELLM.EXTRA_HEADERS": json.dumps({"x-goog-api-key": header_key})}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + handler = LiteLLMAIHandler() + captured, boundaries = [], [] + loop = asyncio.get_running_loop() + run_in_executor = loop.run_in_executor + + def mutate_then_dispatch(executor, function, *args): + if not boundaries: + boundaries.append(function) + if source == "environment": + monkeypatch.setenv("GEMINI_API_KEY", "foreign-key") + monkeypatch.setenv("GEMINI_API_BASE", "https://foreign.example/v1") + else: + monkeypatch.setattr(litellm, "api_key", "foreign-key") + monkeypatch.setattr(litellm, "api_base", "https://foreign.example/v1") + return run_in_executor(executor, function, *args) + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + captured.append(request) + raise TransportReached + + # Mutate only after the adapter's guards; retain native dispatch and URL/auth + # selection and stop at HTTP before response parsing or streaming cleanup. + monkeypatch.setattr(loop, "run_in_executor", mutate_then_dispatch) + monkeypatch.setattr(httpx.AsyncClient, "send", send) + if entrypoint == "stream": + monkeypatch.setattr(handler, "_requires_streaming", lambda model: True) + try: + with pytest.raises(TransportReached): + if entrypoint == "probe": + await handler.probe_completion(f"gemini/{model}") + else: + await handler.chat_completion(f"gemini/{model}", "sys", "usr") + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + assert len(boundaries) == len(captured) == 1 + version = "v1alpha" if "gemini-3" in model else "v1beta" + base = initial_base or f"https://generativelanguage.googleapis.com/{version}" + endpoint = "streamGenerateContent?alt=sse" if entrypoint == "stream" else "generateContent" + assert str(captured[0].url) == f"{base}/models/{model}:{endpoint}" + assert captured[0].headers["x-goog-api-key"] == (header_key or initial_key or DUMMY_LITELLM_API_KEY) + + +@pytest.mark.parametrize(("initial_key", "initial_account", "api_base"), ( + (None, "handler-account", None), + ("handler-key", "handler-account", None), + ("pat/handler-pat", "handler-account", None), + (None, None, None), + ("handler-key", None, None), + ("handler-key", "", None), + ("handler-key", "invalid account", None), + ("handler-key", None, "https://explicit.example"), +)) +@pytest.mark.parametrize("timing", (None, "before_params", "after_params")) +@pytest.mark.parametrize("model_name", ("llama3.1-8b", "claude-3-5-sonnet")) +@pytest.mark.asyncio +async def test_snowflake_transport_keeps_credential_and_account_sources( + monkeypatch, initial_key, initial_account, api_base, timing, model_name, +): + from litellm.litellm_core_utils import logging_worker + + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + if initial_key is not None: + monkeypatch.setenv("SNOWFLAKE_JWT", initial_key) + if initial_account is not None: + monkeypatch.setenv("SNOWFLAKE_ACCOUNT_ID", initial_account) + handler = LiteLLMAIHandler() + if api_base: + # Defensive precedence check: Snowflake has no API base setting today. + handler._provider_request_params["snowflake"] = {"api_base": api_base} + model = f"snowflake/{model_name}" + captured = [] + + class TransportReached(BaseException): + pass + + def send(client, request, **kwargs): + captured.append(request) + raise TransportReached + + async def async_send(client, request, **kwargs): + return send(client, request, **kwargs) + + def mutate(): + monkeypatch.setenv("SNOWFLAKE_JWT", "another-handler-key") + monkeypatch.setenv("SNOWFLAKE_ACCOUNT_ID", "another-account") + + # Keep both native Snowflake endpoint/auth paths; stop only at HTTP transport. + monkeypatch.setattr(httpx.Client, "send", send) + monkeypatch.setattr(httpx.AsyncClient, "send", async_send) + if timing == "before_params": + mutate() + resolve_params = handler._get_provider_request_params + + def request_params(*args, **kwargs): + params = resolve_params(*args, **kwargs) + if timing == "after_params": + mutate() + return params + + monkeypatch.setattr(handler, "_get_provider_request_params", request_params) + valid_request = initial_key and (api_base or initial_account == "handler-account") + try: + if not initial_key: + error, match = ValueError, "Snowflake JWT was not resolved for this request" + elif not api_base and not initial_account: + error, match = ValueError, "Snowflake account was not resolved for this request" + else: + error = TransportReached if valid_request else litellm.APIConnectionError + match = None if valid_request else "Invalid account_id format" + with pytest.raises(error, match=match): + await handler.chat_completion(model=model, system="system", user="test") + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + if not valid_request: + assert not captured + return + assert len(captured) == 1 + endpoint = "messages" if model_name.startswith("claude") else "chat/completions" + base = api_base or "https://handler-account.snowflakecomputing.com" + assert str(captured[0].url) == f"{base}/api/v2/cortex/v1/{endpoint}" + expected_key = initial_key + assert captured[0].headers["authorization"] == f"Bearer {expected_key.removeprefix('pat/')}" + assert captured[0].headers["x-snowflake-authorization-token-type"] == ( + "PROGRAMMATIC_ACCESS_TOKEN" if expected_key.startswith("pat/") else "KEYPAIR_JWT" + ) + assert "account_id" not in json.loads(captured[0].content) + + +@pytest.mark.asyncio +async def test_keyless_openai_placeholder_is_request_local(monkeypatch): + handler = LiteLLMAIHandler() + + kwargs = await _call(handler, "gpt-4o") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + assert litellm.api_key is None + assert litellm.openai_key is None + assert openai.api_key is None + + +def test_keyless_registry_provider_does_not_read_a_missing_environment_name(monkeypatch): + provider_config = type("ProviderConfig", (), {"api_key_env": None})() + registry = MagicMock() + registry.list_providers.return_value = ["keyless"] + registry.get.return_value = provider_config + monkeypatch.setattr(litellm_handler, "JSONProviderRegistry", registry) + + handler = LiteLLMAIHandler() + + assert handler._provider_environment_api_keys == {} + assert litellm_handler._has_live_provider_api_key_environment("keyless") is False + + +@pytest.mark.asyncio +async def test_native_openai_key_is_forwarded_explicitly(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "native-openai-key") + monkeypatch.setattr(litellm, "api_key", "another-request-key") + handler = LiteLLMAIHandler() + + kwargs = await _call(handler, "gpt-4o") + + assert kwargs["api_key"] == "native-openai-key" + assert os.environ["OPENAI_API_KEY"] == "native-openai-key" + assert litellm.api_key == "another-request-key" + + +@pytest.mark.asyncio +async def test_native_openai_like_key_is_not_shadowed_by_placeholder(monkeypatch): + monkeypatch.setenv("OPENAI_LIKE_API_KEY", "native-openai-like-key") + monkeypatch.setattr(litellm, "api_key", "another-request-key") + + kwargs = await _call(LiteLLMAIHandler(), "openai_like/my-model") + + assert kwargs["api_key"] == "native-openai-like-key" + assert litellm.api_key == "another-request-key" + + +@pytest.mark.asyncio +async def test_native_openai_key_does_not_override_openai_like_key(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "native-openai-key") + monkeypatch.setenv("OPENAI_LIKE_API_KEY", "native-openai-like-key") + + kwargs = await _call(LiteLLMAIHandler(), "openai_like/my-model") + + assert kwargs["api_key"] == "native-openai-like-key" + + +@pytest.mark.asyncio +async def test_configured_openai_key_does_not_override_openai_like_key(monkeypatch): + overrides = { + "OPENAI.KEY": "configured-openai-key", + "OPENAI.API_BASE": "https://gateway.example/v1", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + monkeypatch.setenv("OPENAI_LIKE_API_KEY", "native-openai-like-key") + + kwargs = await _call(LiteLLMAIHandler(), "openai_like/my-model") + + assert kwargs["api_base"] == "https://gateway.example/v1" + assert kwargs["api_key"] == "native-openai-like-key" + + +@pytest.mark.asyncio +async def test_configured_openai_key_is_not_sent_to_native_openai_like_endpoint(monkeypatch): + overrides = {"OPENAI.KEY": "configured-openai-key"} + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + monkeypatch.setenv("OPENAI_LIKE_API_BASE", "https://openai-like.example/v1") + + kwargs = await _call(LiteLLMAIHandler(), "openai_like/my-model") + + assert kwargs["api_base"] == "https://openai-like.example/v1" + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.asyncio +async def test_configured_openai_key_reaches_same_native_openai_like_endpoint(monkeypatch): + overrides = { + "OPENAI.KEY": "configured-openai-key", + "OPENAI.API_BASE": "https://gateway.example/v1", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + monkeypatch.setenv("OPENAI_LIKE_API_BASE", "https://gateway.example/v1") + + kwargs = await _call(LiteLLMAIHandler(), "openai_like/my-model") + + assert kwargs["api_base"] == "https://gateway.example/v1" + assert kwargs["api_key"] == "configured-openai-key" + + +@pytest.mark.asyncio +async def test_keyless_openai_like_does_not_receive_placeholder(): + kwargs = await _call(LiteLLMAIHandler(), "openai_like/my-model") + + assert "api_key" not in kwargs + + +@pytest.mark.parametrize("fallback", ("litellm_api_key", "litellm_openai_key", "openai_environment")) +@pytest.mark.asyncio +async def test_openai_like_blocks_unrelated_openai_fallback(monkeypatch, fallback): + if fallback == "litellm_api_key": + monkeypatch.setattr(litellm, "api_key", "unrelated-openai-key") + elif fallback == "litellm_openai_key": + monkeypatch.setattr(litellm, "openai_key", "unrelated-openai-key") + else: + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-openai-key") + + kwargs = await _call(LiteLLMAIHandler(), "openai_like/my-model") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize(("model", "environment_variable"), ( + ("amazon_nova/model", "AMAZON_NOVA_API_KEY"), + ("aleph_alpha/model", "ALEPHALPHA_API_KEY"), + ("anthropic/claude-x", "ANTHROPIC_API_KEY"), + ("azure/gpt-4o", "AZURE_API_KEY"), + ("azure_ai/gpt-4o", "AZURE_AI_API_KEY"), + ("bedrock_mantle/openai.gpt-oss-120b", "AWS_BEARER_TOKEN_BEDROCK"), + ("codestral/codestral-latest", "CODESTRAL_API_KEY"), + ("cohere/command-r", "CO_API_KEY"), + ("cohere_chat/command-r", "COHERE_API_KEY"), + ("cloudflare/model", "CLOUDFLARE_API_KEY"), + ("compactifai/model", "COMPACTIFAI_API_KEY"), + ("dashscope/qwen-max", "DASHSCOPE_API_KEY"), + ("databricks/model", "DATABRICKS_API_KEY"), + ("datarobot/model", "DATAROBOT_API_TOKEN"), + ("deepinfra/model", "DEEPINFRA_API_KEY"), + ("deepseek/deepseek-chat", "DEEPSEEK_API_KEY"), + ("gemini/gemini-2.5-pro", "GEMINI_API_KEY"), + ("gemini/gemini-2.5-pro", "PALM_API_KEY"), + ("gradient_ai/model", "GRADIENT_AI_API_KEY"), + ("groq/model", "GROQ_API_KEY"), + ("heroku/model", "HEROKU_API_KEY"), + ("huggingface/model", "HF_TOKEN"), + ("langflow/model", "LANGFLOW_API_KEY"), + ("langgraph/model", "LANGGRAPH_API_KEY"), + ("mistral/model", "MISTRAL_API_KEY"), + ("moonshot/model", "MOONSHOT_API_KEY"), + ("minimax/model", "MINIMAX_API_KEY"), + ("lemonade/model", "LEMONADE_API_KEY"), + ("ollama/model", "OLLAMA_API_KEY"), + ("openai/gpt-4o", "OPENAI_API_KEY"), + ("openrouter/model", "OPENROUTER_API_KEY"), + ("predibase/model", "PREDIBASE_API_KEY"), + ("replicate/model", "REPLICATE_API_TOKEN"), + ("sambanova/model", "SAMBANOVA_API_KEY"), + ("sap/model", "AICORE_SERVICE_KEY"), + ("text-completion-codestral/codestral-latest", "CODESTRAL_API_KEY"), + ("text-completion-inception/model", "INCEPTION_API_KEY"), + ("veniceai/model", "VENICE_AI_API_KEY"), + ("watsonx/model", "WATSONX_API_KEY"), + ("watsonx_text/model", "WX_API_KEY"), + ("xai/model", "XAI_API_KEY"), + ("xiaomi_mimo/model", "XIAOMI_MIMO_API_KEY"), + ("zai/model", "ZAI_API_KEY"), +)) +@pytest.mark.asyncio +async def test_native_provider_key_shadows_residual_global(monkeypatch, model, environment_variable): + monkeypatch.setenv(environment_variable, "native-provider-key") + monkeypatch.setattr(litellm, "api_key", "another-request-key") + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_key"] == "native-provider-key" + assert litellm.api_key == "another-request-key" + + +@pytest.mark.asyncio +async def test_compatible_provider_key_shadows_residual_global(monkeypatch): + monkeypatch.setenv("TOGETHERAI_API_KEY", "native-together-key") + monkeypatch.setattr(litellm, "api_key", "another-request-key") + resolve_provider = MagicMock(side_effect=AssertionError("explicit providers must use the environment snapshot")) + monkeypatch.setattr(litellm, "get_llm_provider", resolve_provider) + + kwargs = await _call(LiteLLMAIHandler(), "together_ai/model") + + assert kwargs["api_key"] == "native-together-key" + assert litellm.api_key == "another-request-key" + resolve_provider.assert_not_called() + + +@pytest.mark.asyncio +async def test_keyless_compatible_provider_blocks_residual_global(monkeypatch): + monkeypatch.setattr(litellm, "api_key", "another-request-key") + + kwargs = await _call(LiteLLMAIHandler(), "together_ai/model") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + assert litellm.api_key == "another-request-key" + + +@pytest.mark.asyncio +async def test_late_native_provider_key_does_not_bypass_handler_snapshot(monkeypatch): + handler = LiteLLMAIHandler() + monkeypatch.setenv("GROQ_API_KEY", "later-provider-key") + + kwargs = await _call(handler, "groq/model") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.asyncio +async def test_ollama_native_endpoint_is_frozen_with_its_key(monkeypatch): + monkeypatch.setenv("OLLAMA_API_KEY", "native-ollama-key") + monkeypatch.setenv("OLLAMA_API_BASE", "https://tenant-a.example") + handler = LiteLLMAIHandler() + monkeypatch.setenv("OLLAMA_API_BASE", "https://tenant-b.example") + + kwargs = await _call(handler, "ollama/model") + + assert kwargs["api_key"] == "native-ollama-key" + assert kwargs["api_base"] == "https://tenant-a.example" + + +@pytest.mark.asyncio +async def test_moonshot_native_endpoint_is_frozen_with_its_key(monkeypatch): + monkeypatch.setenv("MOONSHOT_API_KEY", "native-moonshot-key") + monkeypatch.setenv("MOONSHOT_API_BASE", "https://tenant-a.example/v1") + handler = LiteLLMAIHandler() + monkeypatch.setenv("MOONSHOT_API_BASE", "https://tenant-b.example/v1") + + kwargs = await _call(handler, "moonshot/model") + + assert kwargs["api_key"] == "native-moonshot-key" + assert kwargs["api_base"] == "https://tenant-a.example/v1" + + +@pytest.mark.asyncio +async def test_late_openai_key_does_not_bypass_openai_like_guard(monkeypatch): + handler = LiteLLMAIHandler() + monkeypatch.setenv("OPENAI_API_KEY", "later-openai-key") + + kwargs = await _call(handler, "openai_like/model") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize(("model", "preferred_variable", "fallback_variable"), ( + ("azure/gpt-4o", "AZURE_OPENAI_API_KEY", "AZURE_API_KEY"), + ("bedrock_mantle/openai.gpt-oss-120b", "BEDROCK_MANTLE_API_KEY", "AWS_BEARER_TOKEN_BEDROCK"), + ("cohere/command-r", "COHERE_API_KEY", "CO_API_KEY"), + ("featherless_ai/model", "FEATHERLESS_AI_API_KEY", "FEATHERLESS_API_KEY"), + ("fireworks_ai/model", "FIREWORKS_API_KEY", "FIREWORKS_AI_API_KEY"), + ("fireworks_ai/model", "FIREWORKS_AI_API_KEY", "FIREWORKSAI_API_KEY"), + ("fireworks_ai/model", "FIREWORKSAI_API_KEY", "FIREWORKS_AI_TOKEN"), + ("friendliai/model", "FRIENDLIAI_API_KEY", "FRIENDLI_TOKEN"), + ("gemini/gemini-2.5-pro", "GOOGLE_API_KEY", "GEMINI_API_KEY"), + ("gemini/gemini-2.5-pro", "GEMINI_API_KEY", "PALM_API_KEY"), + ("huggingface/model", "HF_TOKEN", "HUGGINGFACE_API_KEY"), + ("mistral/model", "MISTRAL_AZURE_API_KEY", "MISTRAL_API_KEY"), + ("openrouter/model", "OPENROUTER_API_KEY", "OR_API_KEY"), + ("perplexity/model", "PERPLEXITYAI_API_KEY", "PERPLEXITY_API_KEY"), + ("replicate/model", "REPLICATE_API_KEY", "REPLICATE_API_TOKEN"), + ("together_ai/model", "TOGETHER_API_KEY", "TOGETHER_AI_API_KEY"), + ("together_ai/model", "TOGETHER_AI_API_KEY", "TOGETHERAI_API_KEY"), + ("together_ai/model", "TOGETHERAI_API_KEY", "TOGETHER_AI_TOKEN"), + ("vercel_ai_gateway/model", "VERCEL_AI_GATEWAY_API_KEY", "VERCEL_OIDC_TOKEN"), + ("watsonx/model", "WATSONX_APIKEY", "WATSONX_API_KEY"), + ("watsonx/model", "WATSONX_API_KEY", "WX_API_KEY"), + ("watsonx_text/model", "WATSONX_APIKEY", "WATSONX_API_KEY"), + ("watsonx_text/model", "WATSONX_API_KEY", "WX_API_KEY"), +)) +@pytest.mark.asyncio +async def test_native_provider_environment_precedence(monkeypatch, model, preferred_variable, fallback_variable): + monkeypatch.setenv(preferred_variable, "preferred-provider-key") + monkeypatch.setenv(fallback_variable, "fallback-provider-key") + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_key"] == "preferred-provider-key" + + +@pytest.mark.parametrize("model", ("watsonx/model", "watsonx_text/model")) +@pytest.mark.asyncio +async def test_watsonx_zen_auth_overrides_generic_api_key(monkeypatch, model): + monkeypatch.setenv("WX_API_KEY", "iam-api-key") + monkeypatch.setenv("WATSONX_ZENAPIKEY", "zen-api-key") + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_key"] == litellm_handler.DUMMY_LITELLM_API_KEY + assert kwargs["zen_api_key"] == "zen-api-key" + + +@pytest.mark.asyncio +async def test_watsonx_text_zen_auth_blocks_live_api_key_fallback(monkeypatch): + monkeypatch.setenv("WATSONX_ZENAPIKEY", "zen-api-key") + + kwargs = await _call(LiteLLMAIHandler(), "watsonx_text/model") + + assert kwargs["api_key"] == litellm_handler.DUMMY_LITELLM_API_KEY + assert kwargs["zen_api_key"] == "zen-api-key" + + +@pytest.mark.asyncio +async def test_request_setting_overrides_native_provider_key(monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "native-anthropic-key") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"ANTHROPIC.KEY": "request-anthropic-key"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), "anthropic/claude-x") + + assert kwargs["api_key"] == "request-anthropic-key" + + +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize(("source", "api_key", "bearer"), ( + ("settings", DUMMY_LITELLM_API_KEY, None), + ("environment", DUMMY_LITELLM_API_KEY, None), + ("settings", DUMMY_LITELLM_API_KEY, "owned-bearer"), + ("environment", DUMMY_LITELLM_API_KEY, "owned-bearer"), + ("settings", "ordinary-key", "owned-bearer"), + ("environment", "ordinary-key", None), + ("keyless", None, None), + ("keyless", None, "owned-bearer"), +)) +@pytest.mark.asyncio +async def test_native_anthropic_captured_literal_key(monkeypatch, entrypoint, source, api_key, bearer): + from litellm.litellm_core_utils import logging_worker + + if entrypoint == "probe": + # The minimal probe is system-only; let native LiteLLM add its user message + # so this test reaches authentication/HTTP without changing probe production code. + monkeypatch.setattr(litellm, "modify_params", True) + overrides = {"ANTHROPIC.KEY": api_key} if source == "settings" else {} + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + if source == "environment": + monkeypatch.setenv("ANTHROPIC_API_KEY", api_key) + if bearer: + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", bearer) + handler = LiteLLMAIHandler() + monkeypatch.setenv("ANTHROPIC_API_KEY", "foreign-key") + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "foreign-bearer") + captured = [] + + async def send(client, request, **kwargs): + captured.append(request) + return httpx.Response(200, request=request, json={ + "id": "test", "type": "message", "role": "assistant", "model": "claude-sonnet-4", + "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }) + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + async def invoke(): + if entrypoint == "chat": + await handler.chat_completion("anthropic/claude-sonnet-4", "sys", "usr") + else: + await handler.probe_completion("anthropic/claude-sonnet-4") + + if api_key is None and bearer is None: + with pytest.raises(litellm.AuthenticationError): + await invoke() + assert not captured + else: + await invoke() + assert len(captured) == 1 + assert str(captured[0].url) == "https://api.anthropic.com/v1/messages" + assert captured[0].headers.get_list("x-api-key") == ([api_key] if api_key else []) + assert captured[0].headers.get_list("authorization") == ([] if api_key else [f"Bearer {bearer}"]) + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + assert litellm_handler._anthropic_request_auth_token.get() is None + assert litellm_handler.AnthropicModelInfo.get_api_key() == "foreign-key" + assert litellm_handler.AnthropicModelInfo.get_auth_token() == "foreign-bearer" + + +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("scenario", ( + "initial", "changed", "removed", "late_only", "missing", "configured", "configured_fallback", "azure_base", +)) +@pytest.mark.asyncio +async def test_native_azure_endpoint_alias_snapshot(monkeypatch, entrypoint, scenario): + from litellm.litellm_core_utils import logging_worker + + endpoint = f"https://owned-{entrypoint}-{scenario.replace('_', '-')}.openai.azure.com" + overrides = {"OPENAI.API_VERSION": "2024-06-01"} + monkeypatch.setenv("AZURE_API_KEY", "owned-key") + if scenario not in ("late_only", "missing"): + monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", endpoint) + if scenario in ("configured", "configured_fallback", "azure_base"): + monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://lower-priority.openai.azure.com") + if scenario == "azure_base": + monkeypatch.setenv("AZURE_API_BASE", endpoint) + else: + overrides["OPENAI.API_BASE"] = endpoint + if scenario == "configured": + overrides["OPENAI.API_TYPE"] = "azure" + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + handler = LiteLLMAIHandler() + if scenario == "removed": + monkeypatch.delenv("AZURE_OPENAI_ENDPOINT") + elif scenario != "initial" and scenario != "missing": + monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://foreign.openai.azure.com") + captured = [] + + async def send(client, request, **kwargs): + captured.append(request) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + async def invoke(): + if entrypoint == "chat": + await handler.chat_completion("azure/gpt-4o", "sys", "usr") + else: + await handler.probe_completion("azure/gpt-4o") + + if scenario == "late_only": + with pytest.raises(ValueError, match="Refusing live api_base environment fallback for provider azure"): + await invoke() + assert not captured + elif scenario == "missing": + with pytest.raises((ValueError, litellm.APIError)): + await invoke() + assert not captured + else: + await invoke() + assert len(captured) == 1 + assert str(captured[0].url) == ( + f"{endpoint}/openai/deployments/gpt-4o/chat/completions?api-version=2024-06-01" + ) + assert captured[0].headers.get_list("api-key") == ["owned-key"] + assert "authorization" not in captured[0].headers + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + + +@pytest.mark.asyncio +async def test_anthropic_auth_token_is_request_local(monkeypatch): + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "request-auth-token") + handler = LiteLLMAIHandler() + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "another-request-token") + + async def completion(**kwargs): + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + assert litellm_handler.AnthropicModelInfo.get_api_key(kwargs["api_key"]) is None + assert litellm_handler.AnthropicModelInfo.get_auth_token("another-request-token") == "request-auth-token" + return _mock_response() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=completion): + await handler.chat_completion(model="anthropic/claude-x", system="sys", user="usr") + + assert litellm_handler.AnthropicModelInfo.get_auth_token() == "another-request-token" + + +@pytest.mark.asyncio +async def test_keyless_anthropic_does_not_borrow_later_environment_credentials(monkeypatch): + handler = LiteLLMAIHandler() + monkeypatch.setenv("ANTHROPIC_API_KEY", "another-request-key") + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "another-request-token") + + async def completion(**kwargs): + assert litellm_handler.AnthropicModelInfo.get_api_key(kwargs["api_key"]) is None + assert litellm_handler.AnthropicModelInfo.get_auth_token("another-request-token") is None + return _mock_response() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=completion): + await handler.chat_completion(model="anthropic/claude-x", system="sys", user="usr") + + +@pytest.mark.asyncio +async def test_concurrent_anthropic_auth_tokens_stay_request_local(monkeypatch): + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "tenant-a-token") + tenant_a = LiteLLMAIHandler() + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "tenant-b-token") + tenant_b = LiteLLMAIHandler() + captured_tokens = {} + + async def completion(**kwargs): + await asyncio.sleep(0) + request_id = kwargs["messages"][-1]["content"] + captured_tokens[request_id] = litellm_handler.AnthropicModelInfo.get_auth_token() + return _mock_response() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=completion): + await asyncio.gather( + tenant_a.chat_completion(model="anthropic/claude-x", system="sys", user="tenant-a"), + tenant_b.chat_completion(model="anthropic/claude-x", system="sys", user="tenant-b"), + ) + + assert captured_tokens == { + "tenant-a": "tenant-a-token", + "tenant-b": "tenant-b-token", + } + + +@pytest.mark.asyncio +async def test_unrelated_native_provider_key_is_not_forwarded(monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "native-anthropic-key") + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize("model", ( + "anthropic/claude-x", + "azure/gpt-4o", + "bedrock_mantle/openai.gpt-oss-120b", + "veniceai/model", +)) +@pytest.mark.asyncio +async def test_provider_without_key_blocks_residual_generic_global(monkeypatch, model): + monkeypatch.setattr(litellm, "api_key", "another-request-key") + if model.startswith("bedrock_mantle/"): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + monkeypatch.setenv("AWS_REGION", "us-east-1") + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize(("provider", "global_name"), ( + (provider, global_name) + for provider, global_names in litellm_handler.PROVIDER_API_KEY_GLOBALS.items() + for global_name in global_names +)) +@pytest.mark.asyncio +async def test_provider_without_key_blocks_residual_provider_global(monkeypatch, provider, global_name): + if provider == "gdc": + monkeypatch.setenv("GDC_API_BASE", "https://owned.example/v1/projects/owned/locations/local/chat/completions") + monkeypatch.setattr(litellm, global_name, "another-request-key") + + kwargs = await _call(LiteLLMAIHandler(), f"{provider}/model") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize(("provider", "global_name"), ( + ("ai21", "ai21_key"), + ("baseten", "baseten_key"), + ("nebius", "nebius_key"), + ("sap", "sap_service_key"), + ("together_ai", "togetherai_api_key"), + ("wandb", "wandb_key"), +)) +def test_provider_global_guard_matrix_is_explicit(provider, global_name): + assert global_name in litellm_handler.PROVIDER_API_KEY_GLOBALS[provider] + + +@pytest.mark.parametrize("model", ("chatgpt/gpt-5", "github_copilot/gpt-4o")) +@pytest.mark.asyncio +async def test_managed_auth_provider_blocks_generic_credentials(monkeypatch, model): + monkeypatch.setattr(litellm, "api_key", "another-request-key") + monkeypatch.setattr(litellm, "openai_key", "another-openai-key") + monkeypatch.setattr(litellm, "headers", {"Authorization": "Bearer another-request-token"}) + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({ + "OPENAI.API_BASE": "https://gateway.example/v1", + "OPENAI.API_VERSION": "request-version", + "OPENAI.KEY": "request-openai-key", + "OPENAI.ORG": "request-organization", + }), + ) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + assert isinstance(kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(kwargs["headers"]["OpenAI-Project"], openai.Omit) + assert not {"api_base", "api_version", "organization"} & kwargs.keys() + + +@pytest.mark.parametrize( + "fallback", + ("litellm_api_key", "litellm_openai_key", "openai_environment", "openai_sdk_key"), +) +@pytest.mark.asyncio +async def test_unknown_provider_blocks_residual_openai_credentials(monkeypatch, fallback): + if fallback == "litellm_api_key": + monkeypatch.setattr(litellm, "api_key", "another-request-key") + elif fallback == "litellm_openai_key": + monkeypatch.setattr(litellm, "openai_key", "another-request-key") + elif fallback == "openai_environment": + monkeypatch.setenv("OPENAI_API_KEY", "another-request-key") + else: + monkeypatch.setattr(openai, "api_key", "another-request-key") + + kwargs = await _call(LiteLLMAIHandler(), "future_provider/model") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.asyncio +async def test_unknown_provider_blocks_openai_environment_added_after_handler_creation(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) + handler = LiteLLMAIHandler() + monkeypatch.setenv("OPENAI_API_KEY", "another-request-key") + + kwargs = await _call(handler, "future_provider/model") + + assert kwargs["api_base"] == "https://gateway.example/v1" + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.asyncio +async def test_keyless_gateway_gets_request_local_placeholder(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), "future_provider/model") + + assert kwargs["api_base"] == "https://gateway.example/v1" + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize("model", ("hosted_vllm/model", "lm_studio/model")) +@pytest.mark.asyncio +async def test_keyless_compatible_gateway_gets_request_local_placeholder(monkeypatch, model): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_base"] == "https://gateway.example/v1" + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize("model", ("hosted_vllm/model", "lm_studio/model")) +@pytest.mark.asyncio +async def test_compatible_gateway_forwards_request_local_openai_environment_key(monkeypatch, model): + monkeypatch.setenv("OPENAI_API_KEY", "gateway-key") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_base"] == "https://gateway.example/v1" + assert kwargs["api_key"] == "gateway-key" + + +@pytest.mark.parametrize(("model", "environment_variable"), ( + ("hosted_vllm/model", "HOSTED_VLLM_API_BASE"), + ("lm_studio/model", "LM_STUDIO_API_BASE"), +)) +@pytest.mark.asyncio +async def test_compatible_native_endpoint_is_frozen_with_keyless_placeholder(monkeypatch, model, environment_variable): + monkeypatch.setenv(environment_variable, "https://tenant-a.example/v1") + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "https://tenant-b.example/v1") + + kwargs = await _call(handler, model) + + assert kwargs["api_base"] == "https://tenant-a.example/v1" + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize(("model", "api_base_variable", "api_key_variable"), ( + ("hosted_vllm/model", "HOSTED_VLLM_API_BASE", "HOSTED_VLLM_API_KEY"), + ("lm_studio/model", "LM_STUDIO_API_BASE", "LM_STUDIO_API_KEY"), +)) +@pytest.mark.asyncio +async def test_compatible_native_endpoint_preserves_native_key( + monkeypatch, + model, + api_base_variable, + api_key_variable, +): + monkeypatch.setenv(api_base_variable, "https://native.example/v1") + monkeypatch.setenv(api_key_variable, "native-key") + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_base"] == "https://native.example/v1" + assert kwargs["api_key"] == "native-key" + + +@pytest.mark.parametrize(("model", "environment_variable"), ( + ("hosted_vllm/model", "HOSTED_VLLM_API_BASE"), + ("lm_studio/model", "LM_STUDIO_API_BASE"), +)) +@pytest.mark.asyncio +async def test_compatible_native_endpoint_overrides_openai_gateway(monkeypatch, model, environment_variable): + monkeypatch.setenv(environment_variable, "https://native.example/v1") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_base"] == "https://native.example/v1" + + +@pytest.mark.parametrize(("model", "environment_variable"), ( + ("hosted_vllm/model", "HOSTED_VLLM_API_BASE"), + ("lm_studio/model", "LM_STUDIO_API_BASE"), +)) +@pytest.mark.asyncio +async def test_late_compatible_native_endpoint_is_rejected(monkeypatch, model, environment_variable): + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "https://another-request.example/v1") + + with pytest.raises(ValueError, match="Refusing live api_base environment fallback"): + await _call(handler, model) + + +@pytest.mark.asyncio +async def test_known_unmapped_provider_blocks_residual_generic_global(monkeypatch): + monkeypatch.setattr(litellm, "api_key", "another-request-key") + + kwargs = await _call(LiteLLMAIHandler(), "vllm/model") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.asyncio +async def test_keyless_custom_openai_placeholder_is_request_local(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), "custom_openai/my-model") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + assert kwargs["api_base"] == "https://gateway.example/v1" + assert litellm.api_key is None + assert litellm.openai_key is None + + +@pytest.mark.asyncio +async def test_custom_openai_key_is_forwarded_without_gateway_base(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.KEY": "openai-key"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), "custom_openai/my-model") + + assert kwargs["api_key"] == "openai-key" + assert kwargs["api_base"] == litellm_handler.OPENAI_DEFAULT_API_BASE + + +@pytest.mark.parametrize("model", ("gpt-4o", "custom_openai/my-model")) +@pytest.mark.parametrize("environment_variable", ("OPENAI_BASE_URL", "OPENAI_API_BASE")) +@pytest.mark.asyncio +async def test_openai_endpoint_added_after_handler_does_not_redirect_request( + monkeypatch, + model, + environment_variable, +): + monkeypatch.setenv("OPENAI_API_KEY", "request-openai-key") + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "https://another-request.example/v1") + + kwargs = await _call(handler, model) + + assert kwargs["api_key"] == "request-openai-key" + assert kwargs["api_base"] == litellm_handler.OPENAI_DEFAULT_API_BASE + + +@pytest.mark.asyncio +async def test_native_custom_openai_key_shadows_residual_global(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "native-openai-key") + monkeypatch.setattr(litellm, "api_key", "another-request-key") + + kwargs = await _call(LiteLLMAIHandler(), "custom_openai/my-model") + + assert kwargs["api_key"] == "native-openai-key" + assert litellm.api_key == "another-request-key" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ("openai_like/my-model", "custom_openai/my-model")) +async def test_openai_gateway_credentials_reach_compatible_aliases(monkeypatch, model): + overrides = { + "OPENAI.KEY": "openai-key", + "OPENAI.API_BASE": "https://gateway.example/v1", + "OPENAI.API_VERSION": "2026-01-01", + "OPENAI.ORG": "openai-org", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["api_base"] == "https://gateway.example/v1" + assert kwargs["api_key"] == "openai-key" + assert "api_version" not in kwargs + assert "organization" not in kwargs + + +@pytest.mark.parametrize("global_name", ("api_key", "openai_key")) +@pytest.mark.asyncio +async def test_residual_litellm_global_key_is_not_reused(monkeypatch, global_name): + monkeypatch.setattr(litellm, global_name, "another-request-key") + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + assert getattr(litellm, global_name) == "another-request-key" + + +@pytest.mark.parametrize( + ("model", "global_name", "global_value", "error"), + ( + ("chatgpt/gpt-5", "api_base", "https://another-request.example/v1", "API base"), + ("github_copilot/gpt-4o", "api_base", "https://another-request.example/v1", "API base"), + ("azure/gpt-4o", "api_version", "another-request-version", "API version"), + ("vertex_ai/gemini-2.5-pro", "vertex_project", "another-request-project", "vertex_project"), + ("vertex_ai/gemini-2.5-pro", "vertex_location", "another-request-location", "vertex_location"), + ), +) +@pytest.mark.asyncio +async def test_residual_litellm_global_routing_is_rejected( + monkeypatch, + model, + global_name, + global_value, + error, +): + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, global_name, global_value) + + with pytest.raises(ValueError, match=error): + await _call(handler, model) + + assert getattr(litellm, global_name) == global_value + + +@pytest.mark.parametrize("provider", sorted(litellm_handler.LITELLM_GLOBAL_FIRST_API_BASE_PROVIDERS)) +def test_request_api_base_does_not_bypass_global_first_litellm_routing_guard(monkeypatch, provider): + monkeypatch.setattr(litellm, "api_base", "https://another-request.example/v1") + + with pytest.raises(ValueError, match="API base"): + litellm_handler._guard_request_routing_globals( + provider, + {"api_base": "https://request.example/v1"}, + ) + + +@pytest.mark.parametrize( + ("model", "environment_variable", "environment_value"), + ( + ("anthropic/claude-x", "ANTHROPIC_API_BASE", "https://another-request.example"), + ("anthropic/claude-x", "ANTHROPIC_BASE_URL", "https://another-request.example"), + ("azure_ai/model", "AZURE_AI_API_BASE", "https://another-request.example"), + ("cohere_chat/model", "COHERE_API_BASE", "https://another-request.example/v2/chat"), + ("databricks/endpoint", "DATABRICKS_API_BASE", "https://another-request.example"), + ("groq/model", "GROQ_API_BASE", "https://another-request.example/v1"), + ("huggingface/model", "HF_API_BASE", "https://another-request.example/v1"), + ("huggingface/model", "HUGGINGFACE_API_BASE", "https://another-request.example/v1"), + ("mistral/model", "MISTRAL_AZURE_API_BASE", "https://another-request.example/v1"), + ("moonshot/model", "MOONSHOT_API_BASE", "https://another-request.example/v1"), + ("ollama/model", "OLLAMA_API_BASE", "https://another-request.example"), + ("openai_like/model", "OPENAI_LIKE_API_BASE", "https://another-request.example/v1"), + ("openrouter/model", "OPENROUTER_API_BASE", "https://another-request.example/v1"), + ("text-completion-codestral/codestral-latest", "CODESTRAL_API_BASE", "https://another-request.example/v1"), + ("azure/gpt-4o", "AZURE_API_BASE", "https://another-request.example"), + ("azure/gpt-4o", "AZURE_API_VERSION", "another-request-version"), + ("vertex_ai/gemini-2.5-pro", "VERTEXAI_PROJECT", "another-request-project"), + ("vertex_ai/gemini-2.5-pro", "GOOGLE_CLOUD_PROJECT", "another-request-project"), + ("vertex_ai/gemini-2.5-pro", "GCLOUD_PROJECT", "another-request-project"), + ("vertex_ai/gemini-2.5-pro", "VERTEXAI_LOCATION", "another-request-location"), + ("vertex_ai/gemini-2.5-pro", "VERTEX_LOCATION", "another-request-location"), + ("vertex_ai/gemini-2.5-pro", "VERTEXAI_API_BASE", "https://another-request.example"), + ("publicai/model", "PUBLICAI_API_BASE", "https://another-request.example/v1"), + ), +) +@pytest.mark.asyncio +async def test_late_provider_routing_environment_is_rejected( + monkeypatch, + model, + environment_variable, + environment_value, +): + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, environment_value) + + with pytest.raises(ValueError, match="environment fallback"): + await _call(handler, model) + + +@pytest.mark.parametrize(("model", "version"), (("gemini-2.5-pro", "v1beta"), ("gemini-3-flash-preview", "v1alpha"))) +@pytest.mark.asyncio +async def test_late_gemini_endpoint_keeps_native_default(monkeypatch, model, version): + handler = LiteLLMAIHandler() + monkeypatch.setenv("GEMINI_API_BASE", "https://another-request.example") + kwargs = await _call(handler, f"gemini/{model}") + assert kwargs["api_base"] == f"https://generativelanguage.googleapis.com/{version}" + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +def test_gemini_default_url_interface_fails_closed(monkeypatch): + from litellm.llms.vertex_ai import common_utils + + monkeypatch.setattr(common_utils, "_get_gemini_url", lambda **kwargs: ("https://unexpected.example", "changed")) + with pytest.raises(RuntimeError, match="Gemini default URL is incompatible"): + LiteLLMAIHandler()._get_provider_request_params("gemini/gemini-2.5-pro") + + +@pytest.mark.parametrize("environment_variable", ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL")) +@pytest.mark.asyncio +async def test_anthropic_endpoint_is_snapshotted_with_its_api_key(monkeypatch, environment_variable): + monkeypatch.setenv("ANTHROPIC_API_KEY", "tenant-a-key") + monkeypatch.setenv(environment_variable, "https://tenant-a.example") + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "https://tenant-b.example") + + kwargs = await _call(handler, "anthropic/claude-x") + + assert kwargs["api_key"] == "tenant-a-key" + assert kwargs["api_base"] == "https://tenant-a.example" + + +@pytest.mark.parametrize( + ("model", "api_key_environment", "api_base_environment"), + ( + ("azure_ai/model", "AZURE_AI_API_KEY", "AZURE_AI_API_BASE"), + ("cohere_chat/model", "COHERE_API_KEY", "COHERE_API_BASE"), + ("gemini/model", "GEMINI_API_KEY", "GEMINI_API_BASE"), + ("groq/model", "GROQ_API_KEY", "GROQ_API_BASE"), + ("huggingface/model", "HF_TOKEN", "HF_API_BASE"), + ("huggingface/model", "HF_TOKEN", "HUGGINGFACE_API_BASE"), + ("mistral/model", "MISTRAL_AZURE_API_KEY", "MISTRAL_AZURE_API_BASE"), + ("mistral/model", "MISTRAL_API_KEY", "MISTRAL_AZURE_API_BASE"), + ("openai_like/model", "OPENAI_LIKE_API_KEY", "OPENAI_LIKE_API_BASE"), + ("text-completion-codestral/codestral-latest", "CODESTRAL_API_KEY", "CODESTRAL_API_BASE"), + ), +) +@pytest.mark.asyncio +async def test_native_provider_endpoint_is_snapshotted_with_its_api_key( + monkeypatch, + model, + api_key_environment, + api_base_environment, +): + monkeypatch.setenv(api_key_environment, "tenant-a-key") + monkeypatch.setenv(api_base_environment, "https://tenant-a.example/v1") + handler = LiteLLMAIHandler() + monkeypatch.setenv(api_base_environment, "https://tenant-b.example/v1") + + kwargs = await _call(handler, model) + + assert kwargs["api_key"] == "tenant-a-key" + assert kwargs["api_base"] == "https://tenant-a.example/v1" + + +@pytest.mark.parametrize( + ("model", "api_base_environment"), + ( + ("github_copilot/gpt-4o", "GITHUB_COPILOT_API_BASE"), + ), +) +@pytest.mark.asyncio +async def test_managed_auth_endpoint_is_snapshotted(monkeypatch, model, api_base_environment): + monkeypatch.setenv(api_base_environment, "https://tenant-a.example/v1") + handler = LiteLLMAIHandler() + monkeypatch.setenv(api_base_environment, "https://tenant-b.example/v1") + + kwargs = await _call(handler, model) + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + assert kwargs["api_base"] == "https://tenant-a.example/v1" + + +@pytest.mark.parametrize("api_base_environment", ("CHATGPT_API_BASE", "OPENAI_CHATGPT_API_BASE")) +@pytest.mark.asyncio +async def test_chatgpt_endpoint_must_remain_unchanged(monkeypatch, api_base_environment): + monkeypatch.setenv(api_base_environment, "https://tenant-a.example/v1") + handler = LiteLLMAIHandler() + + kwargs = await _call(handler, "chatgpt/gpt-5") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + assert kwargs["api_base"] == "https://tenant-a.example/v1" + + +@pytest.mark.parametrize("api_base_environment", ("CHATGPT_API_BASE", "OPENAI_CHATGPT_API_BASE")) +@pytest.mark.parametrize("late_value", ("https://tenant-b.example/v1", None)) +@pytest.mark.asyncio +async def test_changed_chatgpt_endpoint_is_rejected(monkeypatch, api_base_environment, late_value): + monkeypatch.setenv(api_base_environment, "https://tenant-a.example/v1") + handler = LiteLLMAIHandler() + if late_value is None: + monkeypatch.delenv(api_base_environment) + else: + monkeypatch.setenv(api_base_environment, late_value) + + with pytest.raises(ValueError, match="Refusing changed live api_base environment for provider chatgpt"): + await _call(handler, "chatgpt/gpt-5") + + +@pytest.mark.asyncio +async def test_cloudflare_account_is_snapshotted_with_its_api_key(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_API_KEY", "tenant-a-key") + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "tenant-a-account") + handler = LiteLLMAIHandler() + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "tenant-b-account") + + kwargs = await _call(handler, "cloudflare/model") + + assert kwargs["api_key"] == "tenant-a-key" + assert kwargs["api_base"] == "https://api.cloudflare.com/client/v4/accounts/tenant-a-account/ai/v1" + + +@pytest.mark.asyncio +async def test_cloudflare_api_base_is_snapshotted_with_its_api_key(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_API_KEY", "tenant-a-key") + monkeypatch.setenv("CLOUDFLARE_API_BASE", "https://tenant-a.example/v1") + handler = LiteLLMAIHandler() + monkeypatch.setenv("CLOUDFLARE_API_BASE", "https://tenant-b.example/v1") + + kwargs = await _call(handler, "cloudflare/model") + + assert kwargs["api_key"] == "tenant-a-key" + assert kwargs["api_base"] == "https://tenant-a.example/v1" + + +@pytest.mark.asyncio +async def test_late_cloudflare_account_is_rejected(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_API_KEY", "request-key") + handler = LiteLLMAIHandler() + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "another-request-account") + + with pytest.raises(ValueError, match="Refusing live api_base environment fallback"): + await _call(handler, "cloudflare/model") + + +@pytest.mark.parametrize("provider", ("bedrock", "bedrock_mantle")) +@pytest.mark.asyncio +async def test_late_bedrock_runtime_endpoint_is_rejected(monkeypatch, provider): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({ + "aws.AWS_ACCESS_KEY_ID": "request-key", + "aws.AWS_SECRET_ACCESS_KEY": "request-secret", + "aws.AWS_REGION_NAME": "us-east-1", + }), + ) + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", "https://another-request.example") + + with pytest.raises(ValueError, match="Refusing live aws_bedrock_runtime_endpoint environment fallback"): + await _call(handler, f"{provider}/model") + + +@pytest.mark.parametrize("model", ("watsonx/model", "watsonx_text/model")) +@pytest.mark.asyncio +async def test_watsonx_environment_routing_is_snapshotted(monkeypatch, model): + monkeypatch.setenv("WATSONX_APIKEY", "tenant-a-key") + monkeypatch.setenv("WATSONX_URL", "https://tenant-a.example") + monkeypatch.setenv("WATSONX_PROJECT_ID", "tenant-a-project") + monkeypatch.setenv("WATSONX_SPACE_ID", "tenant-a-space") + monkeypatch.setenv("WATSONX_REGION", "tenant-a-region") + handler = LiteLLMAIHandler() + monkeypatch.setenv("WATSONX_URL", "https://tenant-b.example") + monkeypatch.setenv("WATSONX_PROJECT_ID", "tenant-b-project") + monkeypatch.setenv("WATSONX_SPACE_ID", "tenant-b-space") + monkeypatch.setenv("WATSONX_REGION", "tenant-b-region") + + kwargs = await _call(handler, model) + + assert kwargs["api_key"] == "tenant-a-key" + assert kwargs["api_base"] == "https://tenant-a.example" + assert kwargs["project_id"] == "tenant-a-project" + assert kwargs["space_id"] == "tenant-a-space" + assert kwargs["region_name"] == "tenant-a-region" + + +@pytest.mark.parametrize("model", ("watsonx/model", "watsonx_text/model")) +@pytest.mark.parametrize( + ("environment_variable", "parameter"), + ( + ("WATSONX_TOKEN", "token"), + ("WATSONX_ZENAPIKEY", "zen_api_key"), + ), +) +@pytest.mark.asyncio +async def test_watsonx_auth_environment_is_snapshotted(monkeypatch, model, environment_variable, parameter): + monkeypatch.setenv(environment_variable, "tenant-a-credential") + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "tenant-b-credential") + + kwargs = await _call(handler, model) + + if environment_variable == "WATSONX_TOKEN": + assert kwargs["api_key"] == litellm_handler.DUMMY_LITELLM_API_KEY + assert kwargs["headers"]["Authorization"] == "Bearer tenant-a-credential" + assert "token" not in kwargs + else: + assert kwargs["api_key"] == litellm_handler.DUMMY_LITELLM_API_KEY + assert kwargs[parameter] == "tenant-a-credential" + + +@pytest.mark.parametrize("model", ("watsonx/model", "watsonx_text/model")) +@pytest.mark.parametrize("authorization_header", ("Authorization", "authorization")) +@pytest.mark.asyncio +async def test_watsonx_request_header_overrides_snapshotted_token(monkeypatch, model, authorization_header): + monkeypatch.setenv("WATSONX_TOKEN", "environment-token") + monkeypatch.setenv("WATSONX_ZENAPIKEY", "environment-zen-key") + handler = LiteLLMAIHandler() + handler._request_headers = {authorization_header: "Bearer request-token"} + + kwargs = await _call(handler, model) + + assert kwargs["headers"] == {"Authorization": "Bearer request-token"} + assert kwargs["api_key"] == litellm_handler.DUMMY_LITELLM_API_KEY + assert "token" not in kwargs + assert "zen_api_key" not in kwargs + + +@pytest.mark.parametrize("model", ("watsonx/model", "watsonx_text/model")) +@pytest.mark.asyncio +async def test_watsonx_request_authorization_blocks_generic_api_key(monkeypatch, model): + monkeypatch.setenv("WATSONX_APIKEY", "environment-api-key") + handler = LiteLLMAIHandler() + handler._request_headers = {"Authorization": "Bearer request-token"} + + kwargs = await _call(handler, model) + + assert kwargs["headers"] == {"Authorization": "Bearer request-token"} + assert kwargs["api_key"] == litellm_handler.DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize("model", ("watsonx/model", "watsonx_text/model")) +@pytest.mark.asyncio +async def test_watsonx_token_takes_precedence_without_forwarding_zen_key(monkeypatch, model): + monkeypatch.setenv("WATSONX_TOKEN", "request-token") + monkeypatch.setenv("WATSONX_ZENAPIKEY", "request-zen-key") + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["headers"]["Authorization"] == "Bearer request-token" + assert "token" not in kwargs + assert "zen_api_key" not in kwargs + + +@pytest.mark.parametrize("model", ("watsonx/model", "watsonx_text/model")) +@pytest.mark.parametrize( + ("environment_variable", "environment_value"), + ( + ("WATSONX_URL", "https://another-request.example"), + ("WATSONX_PROJECT_ID", "another-request-project"), + ("WATSONX_SPACE_ID", "another-request-space"), + ("WATSONX_REGION", "another-request-region"), + ("WATSONX_TOKEN", "another-request-token"), + ("WATSONX_ZENAPIKEY", "another-request-zen-key"), + ), +) +@pytest.mark.asyncio +async def test_late_watsonx_routing_environment_is_rejected( + monkeypatch, + model, + environment_variable, + environment_value, +): + monkeypatch.setenv("WATSONX_APIKEY", "request-key") + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, environment_value) + + with pytest.raises(ValueError, match="Refusing live .* environment fallback"): + await _call(handler, model) + + +@pytest.mark.asyncio +async def test_bedrock_bearer_token_is_snapshotted(monkeypatch): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "tenant-a-token") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({ + "aws.AWS_ACCESS_KEY_ID": "request-key", + "aws.AWS_SECRET_ACCESS_KEY": "request-secret", + "aws.AWS_REGION_NAME": "us-east-1", + }), + ) + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "tenant-b-token") + + kwargs = await _call(handler, "bedrock/model") + + assert kwargs["api_key"] == "tenant-a-token" + assert kwargs["aws_region_name"] == "us-east-1" + assert "aws_access_key_id" not in kwargs + assert "aws_secret_access_key" not in kwargs + + +@pytest.mark.asyncio +async def test_bedrock_bearer_token_does_not_require_sigv4_credentials(monkeypatch): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "request-bearer-token") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + + kwargs = await _call(LiteLLMAIHandler(), "bedrock/model") + + assert kwargs["api_key"] == "request-bearer-token" + assert kwargs["aws_region_name"] == "us-east-1" + assert "aws_access_key_id" not in kwargs + assert "aws_secret_access_key" not in kwargs + + +@pytest.mark.asyncio +async def test_late_bedrock_bearer_token_is_rejected(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({ + "aws.AWS_ACCESS_KEY_ID": "request-key", + "aws.AWS_SECRET_ACCESS_KEY": "request-secret", + "aws.AWS_REGION_NAME": "us-east-1", + }), + ) + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "another-request-token") + + with pytest.raises(ValueError, match="Refusing process-wide Bedrock bearer token fallback"): + await _call(handler, "bedrock/model") + + +@pytest.mark.asyncio +async def test_late_bedrock_bearer_region_is_rejected(monkeypatch): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "request-bearer-token") + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_DEFAULT_REGION", "another-request-region") + + with pytest.raises(ValueError, match="Refusing live aws_region_name environment fallback"): + await _call(handler, "bedrock/model") + + +@pytest.mark.parametrize("provider", sorted(litellm_handler.AWS_REQUEST_PROVIDERS)) +@pytest.mark.parametrize("selector", litellm_handler.LITELLM_AWS_CREDENTIAL_SELECTOR_ENV_VARS) +@pytest.mark.asyncio +async def test_aws_provider_rejects_ambient_litellm_credential_selector(monkeypatch, provider, selector): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({ + "aws.AWS_ACCESS_KEY_ID": "request-key", + "aws.AWS_SECRET_ACCESS_KEY": "request-secret", + "aws.AWS_REGION_NAME": "us-east-1", + }), + ) + monkeypatch.setenv(selector, "another-request-selector") + + with pytest.raises(ValueError, match=f"Refusing ambient LiteLLM AWS credential selector for provider {provider}"): + await _call(LiteLLMAIHandler(), f"{provider}/model") + + +@pytest.mark.parametrize("provider", ("bedrock", "bedrock_mantle")) +@pytest.mark.asyncio +async def test_bedrock_bearer_does_not_use_ambient_litellm_credential_selector(monkeypatch, provider): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "request-bearer-token") + monkeypatch.setenv("AWS_PROFILE_NAME", "unrelated-profile") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + + kwargs = await _call(LiteLLMAIHandler(), f"{provider}/model") + + assert kwargs["api_key"] == "request-bearer-token" + assert "aws_access_key_id" not in kwargs + assert "aws_secret_access_key" not in kwargs + + +@pytest.mark.asyncio +async def test_bedrock_mantle_environment_routing_is_snapshotted(monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "request-key") + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://tenant-a.example/v1") + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-1") + handler = LiteLLMAIHandler() + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://tenant-b.example/v1") + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-west-2") + + kwargs = await _call(handler, "bedrock_mantle/model") + + assert kwargs["api_base"] == "https://tenant-a.example/v1" + assert kwargs["aws_region_name"] == "us-east-1" + + +@pytest.mark.asyncio +async def test_bedrock_mantle_bearer_uses_configured_aws_region(monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "request-key") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"aws.AWS_REGION_NAME": "ap-northeast-1"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), "bedrock_mantle/model") + + assert kwargs["aws_region_name"] == "ap-northeast-1" + + +@pytest.mark.asyncio +async def test_bedrock_mantle_sigv4_preserves_mantle_region(monkeypatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1") + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-west-2") + + kwargs = await _call(LiteLLMAIHandler(), "bedrock_mantle/model") + + assert kwargs["aws_region_name"] == "us-west-2" + assert kwargs["aws_access_key_id"] == "request-key" + + +@pytest.mark.asyncio +async def test_bedrock_mantle_sigv4_derives_region_from_api_base(monkeypatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-west-2") + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://bedrock-mantle.eu-west-1.api.aws/v1") + + kwargs = await _call(LiteLLMAIHandler(), "bedrock_mantle/model") + + assert kwargs["api_base"] == "https://bedrock-mantle.eu-west-1.api.aws/v1" + assert kwargs["aws_region_name"] == "eu-west-1" + + +@pytest.mark.asyncio +async def test_late_bedrock_mantle_environment_endpoint_is_rejected(monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "request-key") + handler = LiteLLMAIHandler() + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://another-request.example/v1") + + with pytest.raises(ValueError, match="Refusing live .* environment fallback"): + await _call(handler, "bedrock_mantle/model") + + +@pytest.mark.asyncio +async def test_openai_default_endpoint_ignores_residual_litellm_global(monkeypatch): + monkeypatch.setattr(litellm, "api_base", "https://another-request.example/v1") + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert kwargs["api_base"] == litellm_handler.OPENAI_DEFAULT_API_BASE + assert litellm.api_base == "https://another-request.example/v1" + + +@pytest.mark.asyncio +async def test_azure_environment_routing_is_snapshotted(monkeypatch): + monkeypatch.setenv("AZURE_API_BASE", "https://azure-a.example") + monkeypatch.setenv("AZURE_API_VERSION", "tenant-a-version") + handler = LiteLLMAIHandler() + monkeypatch.setenv("AZURE_API_BASE", "https://azure-b.example") + monkeypatch.setenv("AZURE_API_VERSION", "tenant-b-version") + + kwargs = await _call(handler, "azure/gpt-4o") + + assert kwargs["api_base"] == "https://azure-a.example" + assert kwargs["api_version"] == "tenant-a-version" + + +@pytest.mark.parametrize( + ("api_type", "azure_api_base", "expected_api_base", "expected_api_version"), + ( + (None, "https://azure.example", "https://azure.example", "azure-version"), + (None, None, "https://configured.example", "configured-version"), + ("azure", "https://azure.example", "https://configured.example", "configured-version"), + ), +) +@pytest.mark.asyncio +async def test_azure_routing_precedence_depends_on_azure_mode( + monkeypatch, + api_type, + azure_api_base, + expected_api_base, + expected_api_version, +): + overrides = { + "OPENAI.API_BASE": "https://configured.example", + "OPENAI.API_VERSION": "configured-version", + } + if api_type: + overrides["OPENAI.API_TYPE"] = api_type + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + if azure_api_base: + monkeypatch.setenv("AZURE_API_BASE", azure_api_base) + monkeypatch.setenv("AZURE_API_VERSION", "azure-version") + + kwargs = await _call(LiteLLMAIHandler(), "azure/gpt-4o") + + assert kwargs["api_base"] == expected_api_base + assert kwargs["api_version"] == expected_api_version + + +@pytest.mark.asyncio +async def test_azure_environment_base_keeps_environment_version_in_azure_mode(monkeypatch): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.API_VERSION": "configured-version", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + monkeypatch.setenv("AZURE_API_BASE", "https://azure-environment.example") + monkeypatch.setenv("AZURE_API_VERSION", "azure-environment-version") + + kwargs = await _call(LiteLLMAIHandler(), "azure/gpt-4o") + + assert kwargs["api_base"] == "https://azure-environment.example" + assert kwargs["api_version"] == "azure-environment-version" + + +@pytest.mark.asyncio +async def test_azure_ad_endpoint_keeps_configured_api_version(monkeypatch): + overrides = { + "AZURE_AD.CLIENT_ID": "client-id", + "AZURE_AD.API_BASE": "https://azure-ad.example", + "OPENAI.API_VERSION": "configured-version", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_credential", lambda settings: object()) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_token", lambda credential: "azure-ad-token") + monkeypatch.setenv("AZURE_API_BASE", "https://azure-environment.example") + monkeypatch.setenv("AZURE_API_VERSION", "azure-environment-version") + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert kwargs["api_base"] == "https://azure-ad.example" + assert kwargs["api_version"] == "configured-version" + + +@pytest.mark.parametrize("environment_variable", litellm_handler.AZURE_AD_TOKEN_ENV_VARS) +@pytest.mark.asyncio +async def test_azure_ad_token_environment_is_snapshotted(monkeypatch, environment_variable): + monkeypatch.setenv(environment_variable, "tenant-a-token") + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "tenant-b-token") + + kwargs = await _call(handler, "azure/gpt-4o") + + assert kwargs["azure_ad_token"] == "tenant-a-token" + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.parametrize("environment_variable", litellm_handler.AZURE_AD_TOKEN_ENV_VARS) +@pytest.mark.asyncio +async def test_azure_ad_token_added_after_handler_is_rejected(monkeypatch, environment_variable): + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "another-request-token") + + with pytest.raises(ValueError, match="Azure AD token added after handler initialization"): + await _call(handler, "azure/gpt-4o") + + +@pytest.mark.asyncio +async def test_litellm_azure_ad_token_precedes_openai_sdk_environment(monkeypatch): + monkeypatch.setenv("AZURE_AD_TOKEN", "litellm-token") + monkeypatch.setenv("AZURE_OPENAI_AD_TOKEN", "openai-sdk-token") + + kwargs = await _call(LiteLLMAIHandler(), "azure/gpt-4o") + + assert kwargs["azure_ad_token"] == "litellm-token" + + +@pytest.mark.parametrize("location_variable", ("VERTEXAI_LOCATION", "VERTEX_LOCATION")) +@pytest.mark.asyncio +async def test_vertex_environment_routing_is_snapshotted(monkeypatch, location_variable): + monkeypatch.setenv("VERTEXAI_PROJECT", "tenant-a-project") + monkeypatch.setenv(location_variable, "tenant-a-location") + handler = LiteLLMAIHandler() + monkeypatch.setenv("VERTEXAI_PROJECT", "tenant-b-project") + monkeypatch.setenv("VERTEXAI_LOCATION", "tenant-b-location") + + kwargs = await _call(handler, "vertex_ai/gemini-2.5-pro") + + assert kwargs["vertex_project"] == "tenant-a-project" + assert kwargs["vertex_location"] == "tenant-a-location" + + +@pytest.mark.asyncio +async def test_vertex_endpoint_is_snapshotted_with_environment_routing(monkeypatch): + monkeypatch.setenv("VERTEXAI_API_BASE", "https://tenant-a.example/v1") + monkeypatch.setenv("VERTEXAI_PROJECT", "tenant-a-project") + handler = LiteLLMAIHandler() + monkeypatch.setenv("VERTEXAI_API_BASE", "https://tenant-b.example/v1") + monkeypatch.setenv("VERTEXAI_PROJECT", "tenant-b-project") + + kwargs = await _call(handler, "vertex_ai/gemini-2.5-pro") + + assert kwargs["api_base"] == "https://tenant-a.example/v1" + assert kwargs["vertex_project"] == "tenant-a-project" + + +@pytest.mark.asyncio +async def test_vertex_application_credentials_are_snapshotted(monkeypatch, tmp_path): + credentials_path = tmp_path / "vertex-credentials.json" + credentials_path.write_text('{"type": "service_account", "client_email": "tenant-a@example.com"}') + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(credentials_path)) + handler = LiteLLMAIHandler() + credentials_path.write_text('{"type": "service_account", "client_email": "tenant-b@example.com"}') + + kwargs = await _call(handler, "vertex_ai/model") + + assert kwargs["vertex_credentials"] == handler._vertex_gac_adc["cache_key"] + assert handler._vertex_gac_adc["info"]["client_email"] == "tenant-a@example.com" + + +def _impersonated_vertex_adc(tenant, quota_project=None): + return { + "type": "impersonated_service_account", + "service_account_impersonation_url": ( + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/" + f"{tenant}@example.iam.gserviceaccount.com:generateAccessToken" + ), + "source_credentials": { + "type": "authorized_user", + "client_id": f"{tenant}-client", + "client_secret": f"{tenant}-secret", + "refresh_token": f"{tenant}-refresh", + }, + "quota_project_id": quota_project, + } + + +@pytest.mark.parametrize("use_async", (False, True)) +@pytest.mark.parametrize("quota_project", (None, "snapshot-quota")) +@pytest.mark.parametrize("project_environment", ("VERTEXAI_PROJECT", "GOOGLE_CLOUD_PROJECT", "GCLOUD_PROJECT")) +@pytest.mark.parametrize("credential_source", ("GOOGLE_APPLICATION_CREDENTIALS", "VERTEXAI_CREDENTIALS", "inline")) +@pytest.mark.asyncio +async def test_vertex_impersonated_adc_retains_snapshot_through_load_and_refresh( + monkeypatch, tmp_path, use_async, quota_project, project_environment, credential_source, +): + import google.auth + from google.auth import impersonated_credentials + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + monkeypatch.setenv(project_environment, "same-project") + no_default = MagicMock(side_effect=AssertionError("Ambient ADC must not be loaded")) + monkeypatch.setattr(google.auth, "default", no_default) + credentials_path = tmp_path / "impersonated-adc.json" + handlers = [] + snapshots = [] + for tenant in ("tenant-a", "tenant-b"): + snapshot = json.dumps(_impersonated_vertex_adc(tenant, quota_project)) + credentials_path.write_text(snapshot) + if credential_source == "inline": + monkeypatch.setenv("VERTEXAI_CREDENTIALS", snapshot) + else: + monkeypatch.setenv(credential_source, str(credentials_path)) + handler = LiteLLMAIHandler() + handlers.append(handler) + assert json.loads(handler._vertex_credentials) == json.loads(snapshot) + snapshots.append(handler._provider_request_params["vertex_ai"]["vertex_credentials"]) + credentials_path.unlink() + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "late-credentials.json") + monkeypatch.setenv("VERTEXAI_CREDENTIALS", "late-credentials.json") + monkeypatch.setenv("GOOGLE_CLOUD_QUOTA_PROJECT", "late-quota") + for variable in ("VERTEXAI_PROJECT", "GOOGLE_CLOUD_PROJECT", "GCLOUD_PROJECT"): + monkeypatch.setenv(variable, "late-project") + vertex = VertexBase() + refreshed = [] + + def refresh(credentials, request): + assert isinstance(credentials, impersonated_credentials.Credentials) + assert credentials.quota_project_id == quota_project + refreshed.append(credentials.service_account_email) + credentials.token = f"{credentials.service_account_email}:{credentials._source_credentials.refresh_token}" + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1) + + monkeypatch.setattr(impersonated_credentials.Credentials, "refresh", refresh) + observed = [] + + async def completion(**kwargs): + snapshot = kwargs["vertex_credentials"] + if use_async: + token, project = await vertex.get_access_token_async(snapshot, kwargs["vertex_project"]) + else: + token, project = vertex.get_access_token(snapshot, kwargs["vertex_project"]) + observed.append((snapshot, token, project)) + return _mock_response() + + monkeypatch.setattr(litellm_handler, "acompletion", completion) + for probe in (False, True, False): + if probe: + await asyncio.gather(*( + handler.probe_completion("vertex_ai/gemini-2.5-pro", _completion=completion) + for handler in handlers + )) + for credentials, _ in vertex._credentials_project_mapping.values(): + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) - timedelta(hours=1) + else: + await asyncio.gather(*( + handler.chat_completion(model="vertex_ai/gemini-2.5-pro", system="sys", user="usr") + for handler in handlers + )) + for snapshot, token, project in observed: + tenant = "tenant-a" if snapshot == snapshots[0] else "tenant-b" + assert snapshot in snapshots + assert token == f"{tenant}@example.iam.gserviceaccount.com:{tenant}-refresh" + assert project == "same-project" + assert len(observed) == 6 + assert len(vertex._credentials_project_mapping) == 2 + assert len(refreshed) == 4 + no_default.assert_not_called() + assert litellm_handler._vertex_request_credentials.get() is None + + +@pytest.mark.parametrize(("setting", "environment", "expected"), ( + ("configured-project", {"VERTEXAI_PROJECT": "vertex", "GOOGLE_CLOUD_PROJECT": "google", "GCLOUD_PROJECT": "legacy"}, + "configured-project"), + (None, {"VERTEXAI_PROJECT": "vertex", "GOOGLE_CLOUD_PROJECT": "google", "GCLOUD_PROJECT": "legacy"}, "vertex"), + (None, {"GOOGLE_CLOUD_PROJECT": "google", "GCLOUD_PROJECT": "legacy"}, "google"), + (None, {"GCLOUD_PROJECT": "legacy"}, "legacy"), + (None, {"GOOGLE_CLOUD_PROJECT": "", "GCLOUD_PROJECT": "legacy"}, None), +)) +@pytest.mark.asyncio +async def test_vertex_project_snapshot_preserves_explicit_project_precedence(monkeypatch, setting, environment, expected): + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({"VERTEXAI.VERTEX_PROJECT": setting})) + for variable, value in environment.items(): + monkeypatch.setenv(variable, value) + + params = await _call(LiteLLMAIHandler(), "vertex_ai/gemini-2.5-pro") + + assert params.get("vertex_project") == expected + + +@pytest.mark.parametrize("failure", ("source", "project", "refresh")) +@pytest.mark.asyncio +async def test_vertex_impersonated_adc_failures_do_not_fall_back(monkeypatch, failure): + import google.auth + from google.auth.exceptions import InvalidType, RefreshError + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + info = _impersonated_vertex_adc("tenant-a") + if failure == "source": + info["source_credentials"]["type"] = "unsupported" + no_default = MagicMock(side_effect=AssertionError("Ambient ADC must not be loaded")) + monkeypatch.setattr(google.auth, "default", no_default) + vertex = VertexBase() + refresh = MagicMock(side_effect=RefreshError("refresh failed") if failure == "refresh" else None) + monkeypatch.setattr(vertex, "refresh_auth", refresh) + + async def completion(**kwargs): + return vertex.load_auth(kwargs["vertex_credentials"], None if failure == "project" else "project") + + expected = {"source": InvalidType, "project": ValueError, "refresh": RefreshError}[failure] + outer = {"outer": "context"} + token = litellm_handler._vertex_request_credentials.set(outer) + try: + with pytest.raises(expected): + await LiteLLMAIHandler()._acompletion( + _completion=completion, model="vertex_ai/gemini-2.5-pro", vertex_credentials=json.dumps(info), + ) + assert litellm_handler._vertex_request_credentials.get() is outer + finally: + litellm_handler._vertex_request_credentials.reset(token) + no_default.assert_not_called() + + +@pytest.mark.parametrize("context", (None, {"type": "impersonated_service_account", "other": "request"})) +def test_vertex_impersonated_adc_bridge_rejects_unscoped_or_mismatched_input(monkeypatch, context): + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + original = MagicMock(return_value="original") + # A real function avoids MagicMock fabricating the install marker. + def load_original(self, json_obj, scopes): + return original(json_obj, scopes) + + monkeypatch.setattr(VertexBase, "_credentials_from_service_account", load_original) + litellm_handler._install_vertex_impersonated_credentials_bridge() + installed = VertexBase._credentials_from_service_account + litellm_handler._install_vertex_impersonated_credentials_bridge() + assert VertexBase._credentials_from_service_account is installed + info = _impersonated_vertex_adc("tenant-a") + token = litellm_handler._vertex_request_credentials.set(context) + try: + assert VertexBase()._credentials_from_service_account(info, ["scope"]) == "original" + finally: + litellm_handler._vertex_request_credentials.reset(token) + original.assert_called_once_with(info, ["scope"]) + + +@pytest.mark.parametrize("credential_type", ("service_account", "authorized_user", "external_account")) +def test_vertex_impersonated_adc_bridge_preserves_other_loaders(monkeypatch, credential_type): + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + vertex = VertexBase() + credentials = MagicMock(project_id="project") + loader = MagicMock(return_value=credentials) + method = { + "service_account": "_credentials_from_service_account", + "authorized_user": "_credentials_from_authorized_user", + "external_account": "_credentials_from_identity_pool", + }[credential_type] + if credential_type == "service_account": + def load_original(self, json_obj, scopes): + return loader(json_obj, scopes) + monkeypatch.setattr(VertexBase, method, load_original) + else: + monkeypatch.setattr(vertex, method, loader) + monkeypatch.setattr(vertex, "refresh_auth", MagicMock()) + litellm_handler._install_vertex_impersonated_credentials_bridge() + info = {"type": credential_type} + token = litellm_handler._vertex_request_credentials.set(info) + try: + loaded, project = vertex.load_auth(json.dumps(info), "project") + finally: + litellm_handler._vertex_request_credentials.reset(token) + assert loaded is credentials + assert project == "project" + loader.assert_called_once() + + +def _external_vertex_adc(project_number, kind, quota_project): + info = { + "type": "external_account", + "audience": ( + f"//iam.googleapis.com/projects/{project_number}/locations/global/" + "workloadIdentityPools/pool/providers/provider" + ), + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "token_url": "https://sts.googleapis.com/v1/token", + "credential_source": {"file": "/unused-subject-token"}, + } + if quota_project: + info["quota_project_id"] = quota_project + if kind in ("aws", "explicit_aws"): + info["subject_token_type"] = "urn:ietf:params:aws:token-type:aws4_request" + info["credential_source"] = { + "environment_id": "aws1", + "region_url": "http://169.254.169.254/latest/meta-data/placement/availability-zone", + "url": "http://169.254.169.254/latest/meta-data/iam/security-credentials", + "regional_cred_verification_url": "https://sts.{region}.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15", + } + if kind == "explicit_aws": + info.update(aws_access_key_id="snapshot-access", aws_secret_access_key="snapshot-secret", + aws_region_name="us-east-1") + elif kind == "executable": + info["credential_source"] = {"executable": {"command": "never-execute-in-tests"}} + return info + + +@pytest.mark.parametrize("source", ("inline", "vertex_file", "adc_file")) +@pytest.mark.parametrize("explicit_project", (None, "configured-project")) +@pytest.mark.parametrize("entrypoint", ("probe", "completion")) +@pytest.mark.asyncio +async def test_vertex_executable_credentials_fail_before_dispatch( + monkeypatch, tmp_path, source, explicit_project, entrypoint, +): + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({ + "VERTEXAI.VERTEX_PROJECT": explicit_project, + })) + info = json.dumps(_external_vertex_adc("123", "executable", None)) + if source == "inline": + monkeypatch.setenv("VERTEXAI_CREDENTIALS", info) + else: + path = tmp_path / "executable-adc.json" + path.write_text(info) + variable = "VERTEXAI_CREDENTIALS" if source == "vertex_file" else "GOOGLE_APPLICATION_CREDENTIALS" + monkeypatch.setenv(variable, str(path)) + handler = LiteLLMAIHandler() + monkeypatch.setenv("GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES", "1") + completion = AsyncMock(return_value=_mock_response()) + monkeypatch.setattr(litellm_handler, "acompletion", completion) + expected_error = ValueError if entrypoint == "probe" else openai.APIError + with pytest.raises(expected_error, match="executable.*request isolation"): + if entrypoint == "probe": + await handler.probe_completion("vertex_ai/gemini-2.5-pro") + else: + await handler.chat_completion("vertex_ai/gemini-2.5-pro", "sys", "usr") + completion.assert_not_called() + assert litellm_handler._vertex_request_credentials.get() is None + await _call(handler, "gpt-4o") + + +@pytest.mark.parametrize("explicit_project", (None, "configured-project")) +@pytest.mark.asyncio +async def test_vertex_executable_dict_fails_before_loading_credentials(monkeypatch, explicit_project): + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + loader = MagicMock(side_effect=AssertionError("Do not load executable credentials")) + monkeypatch.setattr(VertexBase, "load_auth", loader) + completion = AsyncMock(return_value=_mock_response()) + with pytest.raises(ValueError, match="executable.*request isolation"): + await LiteLLMAIHandler()._acompletion( + _completion=completion, model="vertex_ai/gemini-2.5-pro", vertex_project=explicit_project, + vertex_credentials=_external_vertex_adc("123", "executable", None), + ) + loader.assert_not_called() + completion.assert_not_called() + assert litellm_handler._vertex_request_credentials.get() is None + + +@pytest.mark.parametrize("source", ("cloud_sdk", "direct_path")) +@pytest.mark.parametrize("use_async", (False, True)) +@pytest.mark.parametrize("project", (None, "configured-project")) +@pytest.mark.asyncio +async def test_vertex_executable_default_adc_and_path_are_rejected(monkeypatch, tmp_path, source, use_async, project): + from google.auth import _cloud_sdk, pluggable + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + path = tmp_path / "application_default_credentials.json" + path.write_text(json.dumps(_external_vertex_adc("123", "executable", None))) + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + monkeypatch.setenv("GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES", "1") + monkeypatch.setattr(_cloud_sdk, "get_project_id", lambda: "sdk-project") + helper = MagicMock(side_effect=AssertionError("Executable helper must not run")) + monkeypatch.setattr(pluggable.subprocess, "run", helper) + vertex = VertexBase() + + async def completion(**kwargs): + credentials = kwargs.get("vertex_credentials") + if use_async: + return await vertex.get_access_token_async(credentials, project) + return vertex.get_access_token(credentials, project) + + kwargs = {"vertex_credentials": str(path)} if source == "direct_path" else {} + with pytest.raises(ValueError, match="executable.*request isolation"): + await LiteLLMAIHandler()._acompletion( + _completion=completion, model="vertex_ai/gemini-2.5-pro", **kwargs, + ) + helper.assert_not_called() + assert not vertex._credentials_project_mapping + assert litellm_handler._vertex_request_credentials.get() is None + + +@pytest.mark.parametrize("use_async", (False, True)) +@pytest.mark.parametrize("cache_shape", ("tuple", "legacy")) +@pytest.mark.parametrize("expiry_minutes", (60, 1, -1)) +@pytest.mark.asyncio +async def test_vertex_cached_executable_credentials_are_rejected(monkeypatch, use_async, cache_shape, expiry_minutes): + from google.auth import pluggable + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + credentials = pluggable.Credentials.from_info(_external_vertex_adc("123", "executable", None)) + credentials.token = "previous-identity-token" + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(minutes=expiry_minutes) + vertex = VertexBase() + entry = (credentials, "project") if cache_shape == "tuple" else credentials + vertex._credentials_project_mapping[(None, "project")] = entry + refresh = MagicMock(side_effect=AssertionError("Do not refresh executable credentials")) + monkeypatch.setattr(vertex, "refresh_auth", refresh) + + async def completion(**kwargs): + if use_async: + return await vertex.get_access_token_async(None, "project") + return vertex.get_access_token(None, "project") + + with pytest.raises(ValueError, match="executable.*request isolation"): + await LiteLLMAIHandler()._acompletion(_completion=completion, model="vertex_ai/gemini-2.5-pro") + refresh.assert_not_called() + assert not vertex._background_refresh_tasks + assert vertex._credentials_project_mapping[(None, "project")] is entry + assert litellm_handler._vertex_request_credentials.get() is None + + +@pytest.mark.asyncio +async def test_vertex_executable_guard_delegates_and_restores_context(monkeypatch): + from google.auth import pluggable + from google.auth.credentials import TokenState + + refresh = MagicMock(return_value="original-refresh") + + def original_refresh(credentials, request): + return refresh(credentials, request) + + pluggable.Credentials.refresh = original_refresh + litellm_handler._install_vertex_executable_guard() + descriptors = {name: inspect.getattr_static(pluggable.Credentials, name) + for name in ("from_info", "refresh", "expired", "token_state")} + litellm_handler._install_vertex_executable_guard() + assert all(inspect.getattr_static(pluggable.Credentials, name) is value + for name, value in descriptors.items()) + + class ExecutableCredentials(pluggable.Credentials): + pass + + credentials = ExecutableCredentials.from_info(_external_vertex_adc("123", "executable", None)) + assert isinstance(credentials, ExecutableCredentials) + credentials.token = "outside-request-token" + assert credentials.expired is False + assert credentials.token_state is TokenState.FRESH + request = object() + assert credentials.refresh(request) == "original-refresh" + refresh.assert_called_once_with(credentials, request) + refresh.reset_mock() + + async def completion(**kwargs): + assert litellm_handler._vertex_request_active.get() + with pytest.raises(ValueError, match="executable.*request isolation"): + await asyncio.to_thread(credentials.refresh, request) + raise asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + await LiteLLMAIHandler()._acompletion(_completion=completion, model="vertex_ai/gemini-2.5-pro") + refresh.assert_not_called() + assert not litellm_handler._vertex_request_active.get() + assert litellm_handler._vertex_request_credentials.get() is None + assert credentials.token_state is TokenState.FRESH + + +@pytest.mark.asyncio +async def test_vertex_executable_cache_replacement_after_lock_wait_is_rejected(monkeypatch): + from google.auth import pluggable + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + credentials = pluggable.Credentials.from_info(_external_vertex_adc("123", "executable", None)) + credentials.token = "replacement-identity-token" + vertex = VertexBase() + entered = asyncio.Event() + release = asyncio.Event() + + class WaitingLock: + async def __aenter__(self): + entered.set() + await release.wait() + + async def __aexit__(self, *args): + return False + + monkeypatch.setattr(vertex, "_acquire_async_refresh_lock", lambda key: WaitingLock()) + + async def completion(**kwargs): + return await vertex.get_access_token_async(None, "project") + + task = asyncio.create_task(LiteLLMAIHandler()._acompletion( + _completion=completion, model="vertex_ai/gemini-2.5-pro", + )) + try: + await asyncio.wait_for(entered.wait(), 5) + vertex._credentials_project_mapping[(None, "project")] = (credentials, "project") + release.set() + with pytest.raises(ValueError, match="executable.*request isolation"): + await task + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert not litellm_handler._vertex_request_active.get() + + +@pytest.fixture +def vertex_native_adc_http(monkeypatch): + from urllib.parse import parse_qs + + from google.auth.transport.requests import Request + + calls = [] + + def request(self, url, method="GET", body=None, headers=None, **kwargs): + assert method == "POST" + headers = {key.lower(): value for key, value in (headers or {}).items()} + calls.append((url, headers, body)) + if url.endswith("/fake@example.iam.gserviceaccount.com:generateAccessToken"): + assert headers["authorization"] == "Bearer fake-source-token" + payload = { + "accessToken": "fake-iam-" + headers["x-goog-user-project"], + "expireTime": (datetime.now(UTC) + timedelta(hours=1)).strftime("%Y-%m-%dT%H:%M:%SZ"), + } + else: + assert url in ("https://sts.googleapis.com/v1/token", "https://sts.googleapis.com/v1/oauth/token") + fields = parse_qs(body.decode() if isinstance(body, bytes) else body) + if url.endswith("/oauth/token"): + assert fields["refresh_token"] == ["fake-refresh"] + else: + assert fields["subject_token"] == ["fake-subject-token"] + payload = { + "access_token": "fake-source-token", "expires_in": 3600, "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + } + return type("Response", (), {"status": 200, "data": json.dumps(payload).encode(), "headers": {}})() + + # Keep credential factories, refresh, and header application native; intercept only HTTP. + monkeypatch.setattr(Request, "__call__", request) + return calls + + +@pytest.mark.parametrize("entrypoint", ("chat_completion", "probe_completion")) +@pytest.mark.asyncio +async def test_vertex_gac_workforce_public_native_refresh(monkeypatch, tmp_path, vertex_native_adc_http, entrypoint): + from google.auth import external_account_authorized_user + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + path = tmp_path / "non-sdk-adc.json" + path.write_text(json.dumps({ + "type": "external_account_authorized_user", "client_id": "fake-client", "client_secret": "fake-secret", + "refresh_token": "fake-refresh", "quota_project_id": "file-quota", + "audience": "//iam.googleapis.com/locations/global/workforcePools/fake/providers/fake", + "token_url": "https://sts.googleapis.com/v1/oauth/token", + })) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(path)) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({ + "VERTEXAI.VERTEX_PROJECT": "resource-project", + })) + handler = LiteLLMAIHandler() + vertex = VertexBase() + observed = [] + + async def completion(**kwargs): + observed.append(await vertex.get_access_token_async( + kwargs.get("vertex_credentials"), kwargs.get("vertex_project"), + )) + return _mock_response() + + monkeypatch.setattr(litellm_handler, "acompletion", completion) + for _ in range(2): + if entrypoint == "chat_completion": + assert await handler.chat_completion("vertex_ai/gemini-2.5-pro", "sys", "usr") == ("ok", "stop") + else: + await handler.probe_completion("vertex_ai/gemini-2.5-pro", _completion=completion) + assert observed == [("fake-source-token", "resource-project")] * 2 + assert len(vertex_native_adc_http) == 1 + cached = list(vertex._credentials_project_mapping.values()) + assert cached + assert all( + isinstance(credentials, external_account_authorized_user.Credentials) + and credentials is cached[0][0] and project == "resource-project" + for credentials, project in cached + ) + assert not litellm_handler._vertex_request_active.get() + + +@pytest.mark.parametrize("entrypoint", ("chat_completion", "probe_completion")) +@pytest.mark.parametrize("source", ("gac", "sdk_gac", "vertex")) +@pytest.mark.asyncio +async def test_vertex_adc_public_native_iam_quota_snapshot( + monkeypatch, tmp_path, vertex_native_adc_http, entrypoint, source, +): + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "resource-project") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({ + "VERTEXAI.VERTEX_PROJECT": "resource-project", + })) + subject = tmp_path / "subject-token" + subject.write_text("fake-subject-token") + info = _external_vertex_adc("123", "file", "file-quota") + info["credential_source"] = {"file": str(subject)} + info["service_account_impersonation_url"] = ( + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/" + "fake@example.iam.gserviceaccount.com:generateAccessToken" + ) + path = tmp_path / ("application_default_credentials.json" if source == "sdk_gac" else "non-sdk-adc.json") + path.write_text(json.dumps(info)) + monkeypatch.setenv("VERTEXAI_CREDENTIALS" if source == "vertex" else "GOOGLE_APPLICATION_CREDENTIALS", str(path)) + handlers = [] + for quota in ("initial-quota-a", "initial-quota-b"): + monkeypatch.setenv("GOOGLE_CLOUD_QUOTA_PROJECT", quota) + handlers.append(LiteLLMAIHandler()) + monkeypatch.setenv("GOOGLE_CLOUD_QUOTA_PROJECT", "late-quota") + vertex = VertexBase() + observed = [] + + async def completion(**kwargs): + observed.append(await vertex.get_access_token_async( + kwargs.get("vertex_credentials"), kwargs.get("vertex_project"), + )) + return _mock_response() + + monkeypatch.setattr(litellm_handler, "acompletion", completion) + for _ in range(2): + for handler in handlers: + if entrypoint == "chat_completion": + assert await handler.chat_completion("vertex_ai/gemini-2.5-pro", "sys", "usr") == ("ok", "stop") + else: + await handler.probe_completion("vertex_ai/gemini-2.5-pro", _completion=completion) + expected_quotas = ["file-quota"] if source == "vertex" else ["initial-quota-a", "initial-quota-b"] + iam_calls = [headers for url, headers, _ in vertex_native_adc_http if url.endswith(":generateAccessToken")] + # Assert the actual IAM refresh header, not just a field on the credential object. + assert [headers.get("x-goog-user-project") for headers in iam_calls] == expected_quotas + assert len(vertex_native_adc_http) == 2 * len(expected_quotas) + per_handler_quotas = ["file-quota"] * 2 if source == "vertex" else expected_quotas + assert observed == [("fake-iam-" + quota, "resource-project") for quota in per_handler_quotas] * 2 + cached = list(vertex._credentials_project_mapping.values()) + assert {credentials.quota_project_id for credentials, _ in cached} == set(expected_quotas) + assert len({id(credentials) for credentials, _ in cached}) == len(expected_quotas) + assert all(project == "resource-project" for _, project in cached) + assert not litellm_handler._vertex_request_active.get() + + +@pytest.fixture(scope="module") +def vertex_test_service_account_key(): + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + return key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption(), + ).decode() + + +@pytest.mark.parametrize("credential_type", ("authorized_user", "service_account")) +@pytest.mark.parametrize("source", ("gac", "vertex")) +@pytest.mark.parametrize("quota_override", (None, "", "initial-quota")) +@pytest.mark.asyncio +async def test_vertex_gac_native_factory_and_quota_precedence( + monkeypatch, tmp_path, vertex_test_service_account_key, credential_type, source, quota_override, +): + from google.auth.transport.requests import Request + from google.oauth2 import credentials as user_credentials + from google.oauth2 import service_account + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + info = {"type": credential_type, "quota_project_id": "file-quota"} + if credential_type == "authorized_user": + info.update(client_id="fake-client", client_secret="fake-secret", refresh_token="fake-refresh") + project, credential_class = "resource-project", user_credentials.Credentials + else: + info.update( + client_email="fake@example.iam.gserviceaccount.com", private_key=vertex_test_service_account_key, + token_uri="https://oauth2.googleapis.com/token", project_id="credential-project", + ) + project, credential_class = None, service_account.Credentials + path = tmp_path / "non-sdk-adc.json" + path.write_text(json.dumps(info)) + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS" if source == "gac" else "VERTEXAI_CREDENTIALS", str(path)) + if quota_override is not None: + monkeypatch.setenv("GOOGLE_CLOUD_QUOTA_PROJECT", quota_override) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({"VERTEXAI.VERTEX_PROJECT": project})) + handler = LiteLLMAIHandler() + path.unlink() + monkeypatch.setenv("GOOGLE_CLOUD_QUOTA_PROJECT", "late-quota") + auth_requests = [] + + def request(self, url, method="GET", **kwargs): + assert url == "https://oauth2.googleapis.com/token" and method == "POST" + auth_requests.append(url) + return type("Response", (), { + "status": 200, "headers": {}, + "data": b'{"access_token":"fake-token","expires_in":3600,"token_type":"Bearer"}', + })() + + monkeypatch.setattr(Request, "__call__", request) + vertex = VertexBase() + + async def completion(**kwargs): + assert await vertex.get_access_token_async(kwargs["vertex_credentials"], kwargs.get("vertex_project")) == ( + "fake-token", project or "credential-project", + ) + return _mock_response() + + for _ in range(2): + await handler.probe_completion("vertex_ai/gemini-2.5-pro", _completion=completion) + # Service-account constructors ignore a JSON quota field; only ADC applies + # the nonempty captured environment override. Explicit Vertex does not. + expected_quota = "file-quota" if credential_type == "authorized_user" else None + if source == "gac" and quota_override: + expected_quota = quota_override + assert len(auth_requests) == 1 + assert vertex._credentials_project_mapping + for credentials, resolved_project in vertex._credentials_project_mapping.values(): + assert isinstance(credentials, credential_class) + assert credentials.quota_project_id == expected_quota + assert resolved_project == (project or "credential-project") + + +def _write_vertex_sdk_adc(tmp_path, quota_project): + path = tmp_path / "application_default_credentials.json" + path.write_text(json.dumps({ + "type": "authorized_user", "client_id": "test-client", "client_secret": "test-secret", + "refresh_token": "test-refresh", "quota_project_id": quota_project, + })) + (tmp_path / "configurations").mkdir() + project_config = tmp_path / "configurations" / "config_default" + project_config.write_text("[core]\nproject = sdk-project\n") + return path, project_config + + +@pytest.mark.parametrize("use_async", (False, True)) +@pytest.mark.parametrize("explicit_sdk_path", (False, True)) +@pytest.mark.parametrize("quota_project", (None, "billing-project")) +@pytest.mark.asyncio +async def test_vertex_non_executable_default_adc_preserves_refresh_and_cache( + monkeypatch, tmp_path, use_async, explicit_sdk_path, quota_project, +): + from google.auth import pluggable + from google.oauth2 import credentials as user_credentials + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + path, project_config = _write_vertex_sdk_adc(tmp_path, quota_project) + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + if explicit_sdk_path: + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(path)) + refreshed = [] + + def refresh(credentials, request): + assert credentials.quota_project_id == quota_project + refreshed.append(credentials) + credentials.token = "authorized-user-token" + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1) + + monkeypatch.setattr(user_credentials.Credentials, "refresh", refresh) + vertex = VertexBase() + unrelated = pluggable.Credentials.from_info(_external_vertex_adc("123", "executable", None)) + vertex._credentials_project_mapping[("unrelated", "other-project")] = (unrelated, "other-project") + + async def completion(**kwargs): + if use_async: + return await vertex.get_access_token_async(kwargs["vertex_credentials"], None) + return vertex.get_access_token(kwargs["vertex_credentials"], None) + + handler = LiteLLMAIHandler() + project_config.write_text("[core]\nproject = later-sdk-project\n") + for _ in range(2): + assert await handler._acompletion( + _completion=completion, model="vertex_ai/gemini-2.5-pro", + **handler._provider_request_params["vertex_ai"], + ) == ( + "authorized-user-token", "sdk-project", + ) + assert len(refreshed) == 1 + assert not litellm_handler._vertex_request_active.get() + + +@pytest.mark.parametrize("entrypoint", ("chat_completion", "probe_completion")) +@pytest.mark.parametrize("quota_project", (None, "billing-project")) +@pytest.mark.asyncio +async def test_vertex_explicit_sdk_adc_public_entrypoints_preserve_project_and_cache( + monkeypatch, tmp_path, entrypoint, quota_project, +): + from google.oauth2 import credentials as user_credentials + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + path, project_config = _write_vertex_sdk_adc(tmp_path, quota_project) + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(path)) + handler = LiteLLMAIHandler() + project_config.write_text("[core]\nproject = later-sdk-project\n") + refreshed = [] + + def refresh(credentials, request): + assert credentials.quota_project_id == quota_project + refreshed.append(credentials) + credentials.token = "authorized-user-token" + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1) + + monkeypatch.setattr(user_credentials.Credentials, "refresh", refresh) + vertex = VertexBase() + observed = [] + + async def completion(**kwargs): + observed.append(await vertex.get_access_token_async( + kwargs["vertex_credentials"], kwargs.get("vertex_project"), + )) + return _mock_response() + + monkeypatch.setattr(litellm_handler, "acompletion", completion) + for _ in range(2): + if entrypoint == "chat_completion": + assert await handler.chat_completion( + model="vertex_ai/gemini-2.5-pro", system="sys", user="usr", + ) == ("ok", "stop") + else: + await handler.probe_completion("vertex_ai/gemini-2.5-pro", _completion=completion) + assert observed == [("authorized-user-token", "sdk-project")] * 2 + assert len(refreshed) == 1 + # LiteLLM caches aliases with both the requested and resolved project. + assert vertex._credentials_project_mapping + assert all( + credentials is refreshed[0] and project == "sdk-project" + for credentials, project in vertex._credentials_project_mapping.values() + ) + assert not litellm_handler._vertex_request_active.get() + + +@pytest.mark.parametrize("entrypoint", ("chat_completion", "probe_completion")) +@pytest.mark.parametrize("change", ("late_same_sdk_path", "different_google_path", "late_vertex_credentials")) +@pytest.mark.asyncio +async def test_vertex_explicit_sdk_adc_public_entrypoints_reject_changed_sources( + monkeypatch, tmp_path, entrypoint, change, +): + path, _ = _write_vertex_sdk_adc(tmp_path, "billing-project") + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + if change != "late_same_sdk_path": + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(path)) + handler = LiteLLMAIHandler() + if change == "late_same_sdk_path": + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(path)) + elif change == "different_google_path": + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(tmp_path / "other-adc.json")) + else: + monkeypatch.setenv("VERTEXAI_CREDENTIALS", '{"type": "authorized_user", "client_id": "other-client"}') + completion = AsyncMock(return_value=_mock_response()) + monkeypatch.setattr(litellm_handler, "acompletion", completion) + + try: + with pytest.raises(ValueError, match="Refusing live Vertex credential environment fallback"): + if entrypoint == "chat_completion": + await handler.chat_completion(model="vertex_ai/gemini-2.5-pro", system="sys", user="usr") + else: + await handler.probe_completion("vertex_ai/gemini-2.5-pro", _completion=completion) + finally: + completion.assert_not_called() + + +@pytest.mark.parametrize("selector", ("vertex_inline", "vertex_file", "non_sdk_file", "literal_parent_path")) +@pytest.mark.asyncio +async def test_vertex_explicit_sdk_adc_selector_precedence(monkeypatch, tmp_path, selector): + path, _ = _write_vertex_sdk_adc(tmp_path, "billing-project") + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(path)) + explicit_info = {"type": "authorized_user", "client_id": "explicit-client"} + explicit_file = tmp_path / "explicit-adc.json" + explicit_file.write_text(json.dumps(explicit_info)) + if selector == "vertex_inline": + monkeypatch.setenv("VERTEXAI_CREDENTIALS", json.dumps(explicit_info)) + elif selector == "vertex_file": + monkeypatch.setenv("VERTEXAI_CREDENTIALS", str(explicit_file)) + elif selector == "non_sdk_file": + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(explicit_file)) + else: + # Google Auth compares strings, not canonical paths or file identity. + literal_path = str(tmp_path / "configurations" / ".." / path.name) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", literal_path) + explicit_info = json.loads(path.read_text()) + handler = LiteLLMAIHandler() + kwargs = await _call(handler, "vertex_ai/gemini-2.5-pro") + + if selector in ("non_sdk_file", "literal_parent_path"): + assert kwargs["vertex_credentials"] == handler._vertex_gac_adc["cache_key"] + assert handler._vertex_gac_adc["info"] == explicit_info + assert handler._vertex_gac_adc["sdk_project"] is None + assert handler._vertex_gac_adc["sdk_project_error"] is None + else: + assert json.loads(kwargs["vertex_credentials"]) == explicit_info + assert handler._vertex_gac_adc is None + assert handler._vertex_default_adc is None + + +@pytest.mark.asyncio +async def test_vertex_default_adc_does_not_adopt_another_cloud_sdk_identity(monkeypatch, tmp_path): + from google.auth import _cloud_sdk + from google.oauth2 import credentials as user_credentials + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({ + "VERTEXAI.VERTEX_PROJECT": "shared-project", + })) + monkeypatch.setattr(_cloud_sdk, "get_project_id", lambda: "shared-project") + handlers = [] + for tenant in ("tenant-a", "tenant-b"): + directory = tmp_path / tenant + directory.mkdir() + (directory / "application_default_credentials.json").write_text(json.dumps({ + "type": "authorized_user", "client_id": tenant, "client_secret": "test-secret", + "refresh_token": f"refresh-{tenant}", + })) + monkeypatch.setenv("CLOUDSDK_CONFIG", str(directory)) + handlers.append(LiteLLMAIHandler()) + + def refresh(credentials, request): + credentials.token = f"token-{credentials.client_id}" + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1) + + monkeypatch.setattr(user_credentials.Credentials, "refresh", refresh) + vertex = VertexBase() + observed = [] + + async def completion(**kwargs): + token, _ = await vertex.get_access_token_async( + kwargs.get("vertex_credentials"), kwargs.get("vertex_project"), + ) + observed.append(token) + return _mock_response() + + for _ in range(2): + for handler in handlers: + await handler.probe_completion("vertex_ai/gemini-2.5-pro", _completion=completion) + assert observed == ["token-tenant-a", "token-tenant-b"] * 2 + + +@pytest.mark.parametrize("project_source", ("default", "active_file", "active_environment", "project_environment")) +@pytest.mark.asyncio +async def test_vertex_default_adc_captures_resource_project_separately_from_quota( + monkeypatch, tmp_path, project_source, +): + from google.oauth2 import credentials as user_credentials + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + directory = tmp_path / "sdk" + directory.mkdir() + (directory / "configurations").mkdir() + (directory / "configurations" / "config_default").write_text("[core]\nproject = default-project\n") + (directory / "configurations" / "config_named").write_text("[core]\nproject = named-project\n") + if project_source == "active_file": + (directory / "active_config").write_text("named") + elif project_source == "active_environment": + monkeypatch.setenv("CLOUDSDK_ACTIVE_CONFIG_NAME", "named") + elif project_source == "project_environment": + monkeypatch.setenv("CLOUDSDK_CORE_PROJECT", "environment-project") + path = directory / "application_default_credentials.json" + path.write_text(json.dumps({ + "type": "authorized_user", "client_id": "original-client", "client_secret": "test-secret", + "refresh_token": "test-refresh", "quota_project_id": "file-quota", + })) + monkeypatch.setenv("CLOUDSDK_CONFIG", str(directory)) + monkeypatch.setenv("GOOGLE_CLOUD_QUOTA_PROJECT", "snapshot-quota") + handler = LiteLLMAIHandler() + path.unlink() + (directory / "configurations" / "config_default").write_text("[core]\nproject = late-project\n") + (directory / "configurations" / "config_named").write_text("[core]\nproject = late-project\n") + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path / "other-sdk")) + monkeypatch.setenv("CLOUDSDK_CORE_PROJECT", "late-project") + monkeypatch.setenv("GOOGLE_CLOUD_QUOTA_PROJECT", "late-quota") + + def refresh(credentials, request): + assert credentials.client_id == "original-client" + assert credentials.quota_project_id == "snapshot-quota" + credentials.token = "original-token" + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1) + + monkeypatch.setattr(user_credentials.Credentials, "refresh", refresh) + vertex = VertexBase() + expected_project = { + "default": "default-project", "active_file": "named-project", + "active_environment": "named-project", "project_environment": "environment-project", + }[project_source] + + async def completion(**kwargs): + return await vertex.get_access_token_async(kwargs["vertex_credentials"], None) + + assert await handler._acompletion(_completion=completion, model="vertex_ai/gemini-2.5-pro") == ( + "original-token", expected_project, + ) + + +@pytest.mark.parametrize(("configuration", "legacy", "expected"), ( + (None, None, "named-project"), + ("", None, "named-project"), + ("invalid config", None, "named-project"), + ("NONE", None, "installation-project"), + ("NONE", "[core]\nproject = legacy-project\n", "installation-project"), + (None, "[core]\nproject = legacy-project\n", "legacy-project"), + ("", "[core]\nproject = legacy-project\n", "named-project"), + ("invalid config", "[core]\nproject = legacy-project\n", "legacy-project"), + (None, "# This properties file has been superseded by named configurations.\n" + "# Editing it will have no effect.\n\n[core]\nproject = legacy-project\n", "named-project"), +)) +def test_cloud_sdk_project_snapshot_matches_configuration_precedence( + monkeypatch, tmp_path, configuration, legacy, expected, +): + sdk = tmp_path / "sdk" + sdk.mkdir() + (sdk / ".install").mkdir() + (sdk / "properties").write_text("[core]\nproject = installation-project\n") + monkeypatch.setattr(litellm_handler.shutil, "which", lambda command: str(sdk / "bin" / "gcloud")) + directory = tmp_path / "config" + directory.mkdir() + (directory / "configurations").mkdir() + (directory / "configurations" / "config_default").write_text( + "[core]\nresource = named-project\nproject = %(resource)s\n", + ) + if configuration is not None: + (directory / "active_config").write_text(configuration) + if legacy is not None: + (directory / "properties").write_text(legacy) + assert litellm_handler._snapshot_cloud_sdk_project(str(directory)) == expected + + +@pytest.mark.parametrize(("value", "expected"), ((" example-project ", "example-project"), (" \t ", None))) +def test_cloud_sdk_project_snapshot_normalizes_environment_output(monkeypatch, tmp_path, value, expected): + monkeypatch.setenv("CLOUDSDK_CORE_PROJECT", value) + assert litellm_handler._snapshot_cloud_sdk_project(str(tmp_path)) == expected + + +@pytest.mark.parametrize("source", ("file", "environment")) +@pytest.mark.parametrize(("value", "expected"), ( + ("named", "named-project"), ("named\n", None), ("NONE", None), + (" named ", "default-project"), ("NONE\n", "default-project"), ("\n", "default-project"), +)) +def test_cloud_sdk_project_snapshot_preserves_native_activator_names(monkeypatch, tmp_path, source, value, expected): + monkeypatch.setattr(litellm_handler.shutil, "which", lambda command: None) + (tmp_path / "configurations").mkdir() + (tmp_path / "configurations" / "config_default").write_text("[core]\nproject = default-project\n") + (tmp_path / "configurations" / "config_named").write_text("[core]\nproject = named-project\n") + if source == "file": + (tmp_path / "active_config").write_text(value) + else: + monkeypatch.setenv("CLOUDSDK_ACTIVE_CONFIG_NAME", value) + if value in (" named ", "NONE\n", "\n"): + # Deliberately reject malformed environment selectors rather than + # opening arbitrary paths; the SDK's raw-env behavior is not emulated. + with pytest.raises(ValueError, match="Invalid Cloud SDK configuration name"): + litellm_handler._snapshot_cloud_sdk_project(str(tmp_path)) + return + assert litellm_handler._snapshot_cloud_sdk_project(str(tmp_path)) == expected + if value == "named\n": + # SDK 583 accepts the terminal newline but keeps it in the filename. + (tmp_path / "configurations" / "config_named\n").write_text("[core]\nproject = literal-name-project\n") + assert litellm_handler._snapshot_cloud_sdk_project(str(tmp_path)) == "literal-name-project" + + +@pytest.mark.parametrize("error_source", ("json", "directory")) +@pytest.mark.asyncio +async def test_vertex_default_adc_snapshot_error_cannot_adopt_a_repaired_file(monkeypatch, tmp_path, error_source): + from google.auth import _default + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + path = tmp_path / "application_default_credentials.json" + if error_source == "json": + path.write_text("{") + else: + path.mkdir() + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + handler = LiteLLMAIHandler() + if error_source == "directory": + path.rmdir() + path.write_text(json.dumps({ + "type": "authorized_user", "client_id": "late-client", "client_secret": "late-secret", + "refresh_token": "late-refresh", + })) + discover = MagicMock(side_effect=AssertionError("Must not fall back to managed ADC")) + monkeypatch.setattr(_default, "_get_gce_credentials", discover) + vertex = VertexBase() + + async def completion(**kwargs): + return await vertex.get_access_token_async(kwargs["vertex_credentials"], "project") + + for _ in range(2): + with pytest.raises(ValueError, match="Unable to snapshot default Vertex credentials"): + await handler._acompletion(_completion=completion, model="vertex_ai/gemini-2.5-pro") + discover.assert_not_called() + assert not vertex._credentials_project_mapping + assert litellm_handler._vertex_request_default_adc.get() is None + assert await handler._acompletion(_completion=AsyncMock(return_value="ok"), model="openai/gpt-4o") == "ok" + + +@pytest.mark.parametrize("explicit_project", (None, "explicit-project")) +@pytest.mark.asyncio +async def test_vertex_default_adc_keeps_project_snapshot_error(monkeypatch, tmp_path, explicit_project): + from google.oauth2 import credentials as user_credentials + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + (tmp_path / "application_default_credentials.json").write_text(json.dumps({ + "type": "authorized_user", "client_id": "original-client", "client_secret": "test-secret", + "refresh_token": "test-refresh", + })) + (tmp_path / "configurations").mkdir() + config = tmp_path / "configurations" / "config_default" + config.write_text("invalid configuration") + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + handler = LiteLLMAIHandler() + config.write_text("[core]\nproject = late-project\n") + + def refresh(credentials, request): + credentials.token = "original-token" + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1) + + monkeypatch.setattr(user_credentials.Credentials, "refresh", refresh) + vertex = VertexBase() + + async def completion(**kwargs): + return await vertex.get_access_token_async(kwargs["vertex_credentials"], explicit_project) + + if explicit_project: + assert await handler._acompletion(_completion=completion, model="vertex_ai/gemini-2.5-pro") == ( + "original-token", explicit_project, + ) + else: + with pytest.raises(ValueError, match="Unable to snapshot Cloud SDK project"): + await handler._acompletion(_completion=completion, model="vertex_ai/gemini-2.5-pro") + assert not vertex._credentials_project_mapping + assert litellm_handler._vertex_request_default_adc.get() is None + + +@pytest.mark.asyncio +async def test_vertex_absent_default_adc_does_not_adopt_a_later_file(monkeypatch, tmp_path): + from google.auth import _default + from google.oauth2 import credentials as user_credentials + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + handler = LiteLLMAIHandler() + (tmp_path / "application_default_credentials.json").write_text(json.dumps({ + "type": "authorized_user", "client_id": "late-client", "client_secret": "late-secret", + "refresh_token": "late-refresh", + })) + metadata_credentials = user_credentials.Credentials("managed-token") + metadata_credentials.expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1) + discover = MagicMock(return_value=(metadata_credentials, "managed-project")) + monkeypatch.setattr(_default, "_get_gce_credentials", discover) + monkeypatch.setattr(user_credentials.Credentials, "refresh", lambda self, request: None) + vertex = VertexBase() + + async def completion(**kwargs): + return await vertex.get_access_token_async(kwargs["vertex_credentials"], None) + + assert await handler._acompletion(_completion=completion, model="vertex_ai/gemini-2.5-pro") == ( + "managed-token", "managed-project", + ) + discover.assert_called_once() + assert litellm_handler._vertex_request_default_adc.get() is None + + +@pytest.mark.parametrize("explicit_project", (None, "shared-project")) +@pytest.mark.parametrize("impersonate", (False, True)) +@pytest.mark.parametrize("credential_source", ("explicit", "cloud_sdk", "cloud_sdk_no_project")) +@pytest.mark.asyncio +async def test_vertex_aws_wif_binds_signing_identity_across_cache_and_refresh( + monkeypatch, tmp_path, explicit_project, impersonate, credential_source, +): + from google.auth.transport import requests as google_requests + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({ + "VERTEXAI.VERTEX_PROJECT": explicit_project, + })) + info = _external_vertex_adc("123", "aws", None) + impersonation_url = ( + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/test@example.com:generateAccessToken" + ) + if impersonate: + info["service_account_impersonation_url"] = impersonation_url + if credential_source == "explicit": + monkeypatch.setenv("VERTEXAI_CREDENTIALS", json.dumps(info)) + else: + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + (tmp_path / "application_default_credentials.json").write_text(json.dumps(info)) + if credential_source == "cloud_sdk": + monkeypatch.setenv("CLOUDSDK_CORE_PROJECT", "shared-project") + monkeypatch.setenv("AWS_REGION", "us-east-1") + handlers = [] + for tenant in ("tenant-a", "tenant-b"): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", tenant) + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", f"secret-{tenant}") + monkeypatch.setenv("AWS_SESSION_TOKEN", f"session-{tenant}") + handlers.append(LiteLLMAIHandler()) + + signed_identities = [] + + def request(url, method="GET", headers=None, body=None, **kwargs): + if url == "https://sts.googleapis.com/v1/token": + from urllib.parse import parse_qs + + form = parse_qs(body.decode() if isinstance(body, bytes) else body) + subject = json.loads(unquote(form["subject_token"][0])) + signed_headers = {item["key"].lower(): item["value"] for item in subject["headers"]} + authorization = signed_headers["authorization"] + tenant = authorization.split("Credential=", 1)[1].split("/", 1)[0] + assert signed_headers["x-amz-security-token"] == f"session-{tenant}" + signed_identities.append(tenant) + response = {"access_token": f"token-{tenant}", "expires_in": 3600, "token_type": "Bearer"} + elif url == impersonation_url: + response = { + "accessToken": headers["authorization"].removeprefix("Bearer "), + "expireTime": (datetime.now(UTC) + timedelta(hours=1)).strftime("%Y-%m-%dT%H:%M:%SZ"), + } + else: + assert not explicit_project and credential_source != "cloud_sdk", "Project discovery is unnecessary" + assert url == "https://cloudresourcemanager.googleapis.com/v1/projects/123" + response = {"projectId": "shared-project"} + return type("Response", (), {"status": 200, "data": json.dumps(response).encode()})() + + monkeypatch.setattr(google_requests, "Request", lambda: request) + vertex = VertexBase() + observed = [] + + async def completion(**kwargs): + token, project = await vertex.get_access_token_async( + kwargs["vertex_credentials"], kwargs.get("vertex_project"), + ) + assert project == "shared-project" + observed.append(token) + return _mock_response() + + for round_number in range(3): + if round_number == 2: + for credentials, _ in vertex._credentials_project_mapping.values(): + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) - timedelta(hours=1) + for handler in handlers: + await handler.probe_completion("vertex_ai/gemini-2.5-pro", _completion=completion) + assert observed == ["token-tenant-a", "token-tenant-b"] * 3 + assert set(signed_identities) == {"tenant-a", "tenant-b"} + assert len({id(entry[0]) for entry in vertex._credentials_project_mapping.values()}) == 2 + + +@pytest.mark.parametrize("source", ("explicit_wif", "cloud_sdk", "absent")) +def test_vertex_equivalent_source_snapshots_reuse_cache_identity(monkeypatch, tmp_path, source): + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + info = json.dumps(_external_vertex_adc("123", "aws", None)) + if source == "explicit_wif": + monkeypatch.setenv("VERTEXAI_CREDENTIALS", info) + elif source == "cloud_sdk": + (tmp_path / "application_default_credentials.json").write_text(info) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "snapshot-access") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "snapshot-secret") + monkeypatch.setenv("AWS_REGION", "us-east-1") + snapshots = { + LiteLLMAIHandler()._get_provider_request_params("vertex_ai/gemini-2.5-pro")["vertex_credentials"] + for _ in range(20) + } + assert len(snapshots) == 1 + + +@pytest.mark.parametrize("as_bytes", (False, True)) +def test_vertex_aws_wif_metadata_supplier_keeps_source_and_refresh(monkeypatch, as_bytes): + info = _external_vertex_adc("123", "aws", None) + source = info["credential_source"] + source["imdsv2_session_token_url"] = "http://169.254.169.254/latest/api/token" + source["regional_cred_verification_url"] = "https://sts.{region}.amazonaws.com/custom-verification" + credentials = litellm_handler._vertex_aws_credentials_from_snapshot(info, {}, ["scope-a"]) + credentials = credentials.with_scopes(["scope-b"]).with_quota_project("quota-project") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "late-access") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "late-secret") + monkeypatch.setenv("AWS_REGION", "late-region") + access_key = "metadata-access-a" + requests_seen = [] + + def request(url, method="GET", headers=None, **kwargs): + requests_seen.append((url, method)) + if url == source["imdsv2_session_token_url"]: + assert method == "PUT" + assert headers == {"X-aws-ec2-metadata-token-ttl-seconds": "300"} + body = "metadata-session" + else: + assert headers == {"X-aws-ec2-metadata-token": "metadata-session"} + if url == source["region_url"]: + body = "us-east-1b" + elif url == source["url"]: + body = "bound-role" + else: + assert url == source["url"] + "/bound-role" + body = json.dumps({"AccessKeyId": access_key, "SecretAccessKey": "metadata-secret", "Token": "token"}) + return type("Response", (), {"status": 200, "data": body.encode() if as_bytes else body})() + + for access_key in ("metadata-access-a", "metadata-access-b"): + subject = json.loads(unquote(credentials.retrieve_subject_token(request))) + headers = {item["key"].lower(): item["value"] for item in subject["headers"]} + assert f"Credential={access_key}/" in headers["authorization"] + assert subject["url"] == "https://sts.us-east-1.amazonaws.com/custom-verification" + assert requests_seen.count((source["url"] + "/bound-role", "GET")) == 2 + + +def test_vertex_aws_wif_snapshot_preserves_native_source_validation(): + info = _external_vertex_adc("123", "aws", None) + info["credential_source"]["environment_id"] = "aws2" + with pytest.raises(ValueError, match="aws version"): + litellm_handler._vertex_aws_credentials_from_snapshot(info, {}, ["scope"]) + + +@pytest.mark.parametrize("kind", ("identity_pool", "aws", "explicit_aws")) +@pytest.mark.parametrize("use_async", (False, True)) +@pytest.mark.parametrize("explicit_project", (None, "configured-project")) +@pytest.mark.asyncio +async def test_vertex_wif_snapshot_resolves_project_and_preserves_cache( + monkeypatch, tmp_path, kind, use_async, explicit_project, +): + import google.auth + from google.auth import aws, external_account, identity_pool + from google.auth.transport import requests as google_requests + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({ + "VERTEXAI.VERTEX_PROJECT": explicit_project, + })) + no_default = MagicMock(side_effect=AssertionError("Ambient ADC must not be loaded")) + monkeypatch.setattr(google.auth, "default", no_default) + path = tmp_path / "wif-adc.json" + handlers = [] + snapshots = [] + for number in ("123", "456"): + path.write_text(json.dumps(_external_vertex_adc(number, kind, None))) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(path)) + handler = LiteLLMAIHandler() + handlers.append(handler) + assert handler._vertex_gac_adc["info"] == _external_vertex_adc(number, kind, None) + snapshots.append(handler._provider_request_params["vertex_ai"]["vertex_credentials"]) + path.unlink() + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "late-adc.json") + monkeypatch.setenv("GOOGLE_CLOUD_QUOTA_PROJECT", "late-quota") + vertex = VertexBase() + refreshed = [] + discovered = [] + + def refresh(credentials, request): + assert isinstance(credentials, external_account.Credentials) + assert credentials.quota_project_id is None + refreshed.append(credentials.project_number) + credentials.token = f"token-{credentials.project_number}" + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1) + + def request(url, method, headers, **kwargs): + number = url.rsplit("/", 1)[-1] + assert url == f"https://cloudresourcemanager.googleapis.com/v1/projects/{number}" + assert number in ("123", "456") and method == "GET" + assert headers["authorization"] == f"Bearer token-{number}" + assert "x-goog-user-project" not in headers + discovered.append(number) + return type("Response", (), {"status": 200, "data": json.dumps({"projectId": f"project-{number}"}).encode()})() + + for credentials_type in (identity_pool.Credentials, aws.Credentials): + monkeypatch.setattr(credentials_type, "refresh", refresh) + monkeypatch.setattr(google_requests, "Request", lambda: request) + observed = [] + + async def completion(**kwargs): + snapshot = kwargs["vertex_credentials"] + project = kwargs.get("vertex_project") + if use_async: + token, resolved_project = await vertex.get_access_token_async(snapshot, project) + else: + token, resolved_project = vertex.get_access_token(snapshot, project) + observed.append((snapshot, token, resolved_project)) + return _mock_response() + + for round_number in range(3): + if round_number == 2: + for credentials, _ in vertex._credentials_project_mapping.values(): + credentials.expiry = datetime.now(UTC).replace(tzinfo=None) - timedelta(hours=1) + await asyncio.gather(*( + handler.probe_completion("vertex_ai/gemini-2.5-pro", _completion=completion) + for handler in handlers + )) + for snapshot, token, project in observed: + number = "123" if snapshot == snapshots[0] else "456" + assert snapshot in snapshots + assert token == f"token-{number}" + assert project == (explicit_project or f"project-{number}") + assert len(observed) == 6 + assert sorted(discovered) == ([] if explicit_project else ["123", "456"]) + assert len({id(entry[0]) for entry in vertex._credentials_project_mapping.values()}) == 2 + assert len(refreshed) == (4 if explicit_project else 6) + no_default.assert_not_called() + assert litellm_handler._vertex_request_credentials.get() is None + + +@pytest.mark.parametrize("case", ("no_context", "other_type", "different_snapshot", "invalid_json", "explicit_project")) +def test_vertex_wif_bridge_delegates_unowned_or_explicit_requests(monkeypatch, case): + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + original = MagicMock(return_value=("original-credentials", "original-project")) + + def load_original(self, credentials, project_id): + return original(self, credentials, project_id) + + monkeypatch.setattr(VertexBase, "load_auth", load_original) + info = _external_vertex_adc("123", "identity_pool", None) + context = None if case == "no_context" else dict(info) + credentials = json.dumps(info) + project = "explicit-project" if case == "explicit_project" else None + if case == "other_type": + context["type"] = "service_account" + elif case == "different_snapshot": + context["audience"] = "different-request" + elif case == "invalid_json": + credentials = "not-json" + litellm_handler._install_vertex_wif_project_bridge() + bridge = VertexBase.load_auth + litellm_handler._install_vertex_wif_project_bridge() + assert VertexBase.load_auth is bridge + vertex = VertexBase() + token = litellm_handler._vertex_request_credentials.set(context) + try: + assert vertex.load_auth(credentials, project) == ("original-credentials", "original-project") + finally: + litellm_handler._vertex_request_credentials.reset(token) + original.assert_called_once_with(vertex, credentials, project) + + +@pytest.mark.parametrize("failure", ("missing", "wrong_type", "discovery", "refresh")) +@pytest.mark.asyncio +async def test_vertex_wif_failure_does_not_cache_or_fallback(monkeypatch, failure): + import google.auth + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + info = _external_vertex_adc("123", "identity_pool", None) + credentials = MagicMock() + credentials.get_project_id.return_value = None if failure == "missing" else 123 if failure == "wrong_type" else "project" + if failure == "discovery": + credentials.get_project_id.side_effect = RuntimeError("discovery failed") + vertex = VertexBase() + monkeypatch.setattr(vertex, "_credentials_from_identity_pool", MagicMock(return_value=credentials)) + refresh = MagicMock(side_effect=RuntimeError("refresh failed") if failure == "refresh" else None) + monkeypatch.setattr(vertex, "refresh_auth", refresh) + no_default = MagicMock(side_effect=AssertionError("Ambient ADC must not be loaded")) + monkeypatch.setattr(google.auth, "default", no_default) + + async def completion(**kwargs): + return await vertex.get_access_token_async(kwargs["vertex_credentials"], None) + + expected = {"missing": ValueError, "wrong_type": TypeError, "discovery": RuntimeError, "refresh": RuntimeError}[failure] + with pytest.raises(expected): + await LiteLLMAIHandler()._acompletion( + _completion=completion, model="vertex_ai/gemini-2.5-pro", vertex_credentials=json.dumps(info), + ) + assert not vertex._credentials_project_mapping + no_default.assert_not_called() + assert litellm_handler._vertex_request_credentials.get() is None + + +@pytest.mark.asyncio +async def test_vertex_wif_cancelled_load_keeps_worker_context_isolated(monkeypatch): + import threading + + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + entered = threading.Event() + release = threading.Event() + finished = threading.Event() + info = _external_vertex_adc("123", "identity_pool", None) + credentials = MagicMock() + observed_context = [] + + def discover(request): + entered.set() + assert release.wait(5) + observed_context.append(litellm_handler._vertex_request_credentials.get()) + return "project" + + credentials.get_project_id.side_effect = discover + vertex = VertexBase() + monkeypatch.setattr(vertex, "_credentials_from_identity_pool", MagicMock(return_value=credentials)) + monkeypatch.setattr(vertex, "refresh_auth", MagicMock()) + + def load_and_signal(credentials, project_id): + try: + return VertexBase.load_auth(vertex, credentials, project_id) + finally: + finished.set() + + monkeypatch.setattr(vertex, "load_auth", load_and_signal) + + async def completion(**kwargs): + return await vertex.get_access_token_async(kwargs["vertex_credentials"], None) + + task = asyncio.create_task(LiteLLMAIHandler()._acompletion( + _completion=completion, model="vertex_ai/gemini-2.5-pro", vertex_credentials=json.dumps(info), + )) + try: + assert await asyncio.to_thread(entered.wait, 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert litellm_handler._vertex_request_credentials.get() is None + assert not vertex._credentials_project_mapping + finally: + release.set() + assert await asyncio.to_thread(finished.wait, 5) + assert observed_context == [info] + + +@pytest.mark.parametrize("environment_variable", ("VERTEXAI_CREDENTIALS", "GOOGLE_APPLICATION_CREDENTIALS")) +@pytest.mark.asyncio +async def test_late_vertex_credentials_are_rejected(monkeypatch, environment_variable): + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "another-request-credentials") + + with pytest.raises(ValueError, match="Refusing live Vertex credential environment fallback"): + await _call(handler, "vertex_ai/model") + + +@pytest.mark.parametrize("environment_variable", ("VERTEXAI_CREDENTIALS", "GOOGLE_APPLICATION_CREDENTIALS")) +@pytest.mark.asyncio +async def test_stale_vertex_credentials_only_fail_vertex_requests(monkeypatch, tmp_path, environment_variable): + credentials_path = tmp_path / "missing-vertex-credentials.json" + monkeypatch.setenv(environment_variable, str(credentials_path)) + handler = LiteLLMAIHandler() + + await _call(handler, "anthropic/claude-x") + with pytest.raises(ValueError, match="Unable to snapshot explicit Vertex credentials: FileNotFoundError"): + await _call(handler, "vertex_ai/model") + + +@pytest.mark.asyncio +async def test_openai_sdk_global_does_not_suppress_litellm_placeholder(monkeypatch): + monkeypatch.setattr(openai, "api_key", "openai-sdk-only-key") + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.asyncio +async def test_openai_request_settings_are_forwarded(monkeypatch): + overrides = { + "OPENAI.KEY": "request-openai-key", + "OPENAI.API_BASE": "https://openai.example/v1", + "OPENAI.API_VERSION": "2026-01-01", + "OPENAI.ORG": "org-request", + "LITELLM.EXTRA_HEADERS": ( + '{"openai-organization": "extra-org", "openai-project": "project-request"}' + ), + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert kwargs["api_key"] == "request-openai-key" + assert kwargs["api_base"] == "https://openai.example/v1" + assert kwargs["api_version"] == "2026-01-01" + assert kwargs["organization"] == "org-request" + assert kwargs["headers"]["OpenAI-Organization"] == "org-request" + assert "openai-organization" not in kwargs["headers"] + assert kwargs["headers"]["OpenAI-Project"] == "project-request" + assert "openai-project" not in kwargs["headers"] + + +@pytest.mark.parametrize( + "model", + ("gpt-4o", "azure/gpt-4o", "deepinfra/model"), +) +@pytest.mark.asyncio +async def test_openai_sdk_request_blocks_residual_organization_and_headers(monkeypatch, model): + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "organization", "another-request-organization") + monkeypatch.setattr(litellm, "headers", {"X-Another-Request": "secret"}) + monkeypatch.setenv("OPENAI_PROJECT_ID", "another-request-project") + + kwargs = await _call(handler, model) + + assert isinstance(kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(kwargs["headers"]["OpenAI-Project"], openai.Omit) + assert "X-Another-Request" not in kwargs["headers"] + assert litellm.organization == "another-request-organization" + assert litellm.headers == {"X-Another-Request": "secret"} + + +@pytest.mark.parametrize( + "model", + ( + "aiohttp_openai/gpt-4o", + "openai_like/model", + "openrouter/model", + "deepseek/model", + "groq/model", + "gpt-5-pro", + "azure/codex-mini", + "azure/responses/custom-deployment", + "openai/responses/gpt-4o", + "chatgpt/gpt-5.2", + "github/gpt-5-pro", + "litellm_proxy/gpt-5-pro", + "perplexity/openai/gpt-5.2", + "perplexity/perplexity/glm-5.2", + ), +) +@pytest.mark.asyncio +async def test_raw_http_request_rejects_residual_headers_without_omit_values(monkeypatch, model): + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "headers", {"X-Another-Request": "secret"}) + monkeypatch.setenv("OPENAI_PROJECT_ID", "another-request-project") + + with pytest.raises(ValueError, match="Refusing process-wide LiteLLM headers fallback"): + await _call(handler, model) + + +@pytest.mark.parametrize( + "model", + ( + "aiohttp_openai/gpt-4o", + "openai_like/model", + "openrouter/model", + "deepseek/model", + "groq/model", + "gpt-5-pro", + "azure/codex-mini", + "azure/responses/custom-deployment", + "openai/responses/gpt-4o", + "chatgpt/gpt-5.2", + "github/gpt-5-pro", + "litellm_proxy/gpt-5-pro", + "perplexity/openai/gpt-5.2", + "perplexity/perplexity/glm-5.2", + ), +) +@pytest.mark.asyncio +async def test_raw_http_request_does_not_forward_omit_headers(model): + kwargs = await _call(LiteLLMAIHandler(), model) + + assert "headers" not in kwargs + + +@pytest.mark.asyncio +async def test_custom_aiohttp_openai_provider_does_not_forward_omit_headers(monkeypatch): + settings = _make_settings() + settings.litellm.custom_llm_provider = "aiohttp_openai" + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert kwargs["custom_llm_provider"] == "aiohttp_openai" + assert "headers" not in kwargs + + +@pytest.mark.parametrize("model", ("huggingface/model", "ollama/model", "lemonade/model")) +@pytest.mark.asyncio +async def test_non_openai_request_does_not_inspect_responses_model_info(monkeypatch, model): + get_model_info = MagicMock(side_effect=AssertionError("unexpected model-info lookup")) + monkeypatch.setattr(litellm_handler, "_get_model_info_helper", get_model_info) + + await _call(LiteLLMAIHandler(), model) + + get_model_info.assert_not_called() + + +@pytest.mark.parametrize("model", ("gpt-4o", "cerebras/model")) +@pytest.mark.asyncio +async def test_experimental_openai_raw_http_handler_does_not_forward_omit_headers(monkeypatch, model): + monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert "headers" not in kwargs + + +@pytest.mark.parametrize("value", ("1", "yes")) +@pytest.mark.asyncio +async def test_unrecognized_experimental_handler_values_keep_sdk_omit_headers(monkeypatch, value): + monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", value) + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert isinstance(kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(kwargs["headers"]["OpenAI-Project"], openai.Omit) + + +@pytest.mark.asyncio +async def test_text_completion_openai_prefix_ignores_experimental_chat_handler(monkeypatch): + monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") + monkeypatch.setattr(litellm, "headers", {"X-Another-Request": "secret"}) + + kwargs = await _call(LiteLLMAIHandler(), "text-completion-openai/gpt-3.5-turbo-instruct") + + assert isinstance(kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(kwargs["headers"]["OpenAI-Project"], openai.Omit) + assert "X-Another-Request" not in kwargs["headers"] + + +@pytest.mark.asyncio +async def test_bare_text_completion_model_ignores_experimental_chat_handler(monkeypatch): + monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") + monkeypatch.setattr(litellm, "headers", {"X-Another-Request": "secret"}) + + kwargs = await _call(LiteLLMAIHandler(), "gpt-3.5-turbo-instruct") + + assert isinstance(kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(kwargs["headers"]["OpenAI-Project"], openai.Omit) + assert "X-Another-Request" not in kwargs["headers"] + + +@pytest.mark.parametrize("model", ("ft:babbage-002:acme::abc", "ft:davinci-002:acme::abc")) +@pytest.mark.asyncio +async def test_fine_tuned_text_completion_model_ignores_experimental_chat_handler(monkeypatch, model): + monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") + monkeypatch.setattr(litellm, "headers", {"X-Another-Request": "secret"}) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert isinstance(kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(kwargs["headers"]["OpenAI-Project"], openai.Omit) + assert "X-Another-Request" not in kwargs["headers"] + + +@pytest.mark.asyncio +async def test_route_all_openai_to_responses_does_not_forward_omit_headers(monkeypatch): + monkeypatch.setattr(litellm, "route_all_chat_openai_to_responses", True) + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert "headers" not in kwargs + + +@pytest.mark.asyncio +async def test_route_all_openai_to_responses_keeps_text_completion_sdk_headers(monkeypatch): + monkeypatch.setattr(litellm, "route_all_chat_openai_to_responses", True) + + kwargs = await _call(LiteLLMAIHandler(), "gpt-3.5-turbo-instruct") + + assert isinstance(kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(kwargs["headers"]["OpenAI-Project"], openai.Omit) + + +@pytest.mark.asyncio +async def test_openai_responses_request_forwards_only_explicit_organization(monkeypatch): + overrides = {"OPENAI.ORG": "org-request"} + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + monkeypatch.setattr(litellm, "headers", {"X-Another-Request": "secret"}) + + kwargs = await _call(LiteLLMAIHandler(), "gpt-5-pro") + + assert kwargs["headers"] == {"OpenAI-Organization": "org-request"} + + +@pytest.mark.asyncio +async def test_openai_responses_request_rejects_residual_organization(monkeypatch): + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "organization", "another-request-organization") + + with pytest.raises(ValueError, match="Refusing process-wide LiteLLM organization fallback"): + await _call(handler, "gpt-5-pro") + + +@pytest.mark.asyncio +async def test_openai_responses_request_snapshots_environment_organization(monkeypatch): + monkeypatch.setenv("OPENAI_ORGANIZATION", "request-organization") + handler = LiteLLMAIHandler() + monkeypatch.setenv("OPENAI_ORGANIZATION", "another-request-organization") + + kwargs = await _call(handler, "gpt-5-pro") + + assert kwargs["organization"] == "request-organization" + assert kwargs["headers"] == {"OpenAI-Organization": "request-organization"} + + +@pytest.mark.asyncio +async def test_openai_responses_request_rejects_late_environment_organization(monkeypatch): + handler = LiteLLMAIHandler() + monkeypatch.setenv("OPENAI_ORGANIZATION", "another-request-organization") + + with pytest.raises(ValueError, match="Refusing live organization environment fallback"): + await _call(handler, "gpt-5-pro") + + +@pytest.mark.asyncio +async def test_unmapped_azure_responses_model_keeps_sdk_headers(): + kwargs = await _call(LiteLLMAIHandler(), "azure/codex-mini-latest") + + assert isinstance(kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(kwargs["headers"]["OpenAI-Project"], openai.Omit) + + +@pytest.mark.asyncio +async def test_azure_deployment_id_selects_sdk_header_guard(monkeypatch): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.DEPLOYMENT_ID": "custom-deployment", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + monkeypatch.setattr(litellm, "headers", {"X-Another-Request": "secret"}) + + kwargs = await _call(LiteLLMAIHandler(), "gpt-5-pro") + + assert kwargs["deployment_id"] == "custom-deployment" + assert isinstance(kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(kwargs["headers"]["OpenAI-Project"], openai.Omit) + assert "X-Another-Request" not in kwargs["headers"] + + +@pytest.mark.asyncio +async def test_azure_deployment_id_selects_responses_header_guard(monkeypatch): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.DEPLOYMENT_ID": "gpt-5-pro", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "custom-deployment") + + assert kwargs["deployment_id"] == "gpt-5-pro" + assert "headers" not in kwargs + + +@pytest.mark.asyncio +async def test_azure_deployment_id_matches_responses_model_case_insensitively(monkeypatch): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.DEPLOYMENT_ID": "GPT-5-Pro", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "custom-deployment") + + assert kwargs["deployment_id"] == "GPT-5-Pro" + assert "headers" not in kwargs + + +@pytest.mark.asyncio +async def test_stacked_azure_openai_gpt5_prefix_selects_responses_header_guard(): + kwargs = await _call(LiteLLMAIHandler(), "azure/openai/gpt-5-pro") + + assert kwargs["model"] == "azure/gpt-5-pro" + assert "headers" not in kwargs + + +@pytest.mark.asyncio +async def test_openai_request_merges_explicit_extra_headers(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({ + "LITELLM.EXTRA_HEADERS": ( + '{"openai-organization": "request-organization", ' + '"OPENAI-PROJECT": "request-project", "X-Request": "request-value"}' + ), + }), + ) + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert kwargs["headers"]["OpenAI-Organization"] == "request-organization" + assert kwargs["headers"]["OpenAI-Project"] == "request-project" + assert "openai-organization" not in kwargs["headers"] + assert "OPENAI-PROJECT" not in kwargs["headers"] + assert kwargs["headers"]["X-Request"] == "request-value" + assert "extra_headers" not in kwargs + + +@pytest.mark.parametrize("model", ("anthropic/claude-x", "cohere/command-r")) +@pytest.mark.asyncio +async def test_non_openai_request_rejects_residual_global_headers(monkeypatch, model): + monkeypatch.setattr(litellm, "headers", {"Authorization": "Bearer another-request-token"}) + + with pytest.raises(ValueError, match="Refusing process-wide LiteLLM headers fallback"): + await _call(LiteLLMAIHandler(), model) + + +@pytest.mark.asyncio +async def test_non_openai_request_uses_snapshotted_extra_headers(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"LITELLM.EXTRA_HEADERS": '{"X-Request": "request-value"}'}), + ) + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "headers", {"Authorization": "Bearer another-request-token"}) + + kwargs = await _call(handler, "anthropic/claude-x") + + assert kwargs["headers"] == {"X-Request": "request-value"} + assert litellm.headers == {"Authorization": "Bearer another-request-token"} + + +@pytest.mark.parametrize("key_source", ("settings", "environment")) +@pytest.mark.asyncio +async def test_openai_gateway_base_fallback_forwards_request_local_credentials(monkeypatch, key_source): + overrides = { + "OPENAI.API_BASE": "https://gateway.example/v1", + "OPENAI.API_VERSION": "2026-01-01", + "OPENAI.ORG": "openai-org", + } + if key_source == "settings": + overrides["OPENAI.KEY"] = "gateway-key" + else: + monkeypatch.setenv("OPENAI_API_KEY", "gateway-key") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "meta-llama/model") + + assert kwargs["api_base"] == "https://gateway.example/v1" + assert kwargs["api_key"] == "gateway-key" + assert "api_version" not in kwargs + assert "organization" not in kwargs + + +@pytest.mark.asyncio +async def test_openai_gateway_base_does_not_redirect_native_provider_credentials(monkeypatch): + overrides = { + "OPENAI.KEY": "gateway-key", + "OPENAI.API_BASE": "https://gateway.example/v1", + "OPENAI.API_VERSION": "2026-01-01", + "OPENAI.ORG": "openai-org", + "ANTHROPIC.KEY": "anthropic-key", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "anthropic/claude-x") + + assert "api_base" not in kwargs + assert kwargs["api_key"] == "anthropic-key" + assert "api_version" not in kwargs + assert "organization" not in kwargs + + +@pytest.mark.parametrize( + ("model", "overrides", "expected_params"), + ( + ( + "bedrock/anthropic.claude-v2", + { + "aws.AWS_ACCESS_KEY_ID": "request-access-key", + "aws.AWS_SECRET_ACCESS_KEY": "request-secret-key", + "aws.AWS_REGION_NAME": "us-east-1", + }, + { + "aws_access_key_id": "request-access-key", + "aws_secret_access_key": "request-secret-key", + "aws_region_name": "us-east-1", + }, + ), + ( + "vertex_ai/gemini-2.5-pro", + { + "VERTEXAI.VERTEX_PROJECT": "request-project", + "VERTEXAI.VERTEX_LOCATION": "us-central1", + }, + { + "vertex_project": "request-project", + "vertex_location": "us-central1", + }, + ), + ), +) +@pytest.mark.asyncio +async def test_openai_gateway_base_does_not_redirect_native_provider_context( + monkeypatch, + model, + overrides, + expected_params, +): + overrides = {**overrides, "OPENAI.API_BASE": "https://gateway.example/v1"} + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert "api_base" not in kwargs + assert expected_params.items() <= kwargs.items() + + +@pytest.mark.asyncio +async def test_openai_gateway_credentials_reach_compatible_provider(monkeypatch): + overrides = { + "OPENAI.KEY": "gateway-key", + "OPENAI.API_BASE": "https://gateway.example/v1", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "together_ai/model") + + assert kwargs["api_key"] == "gateway-key" + assert kwargs["api_base"] == "https://gateway.example/v1" + + +@pytest.mark.asyncio +async def test_mosaico_gateway_credentials_reach_openrouter(monkeypatch): + overrides = { + "OPENAI.KEY": "gateway-key", + "OPENAI.API_BASE": "https://openrouter.ai/api/v1", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "openrouter/mistralai/devstral-small") + + assert kwargs["api_key"] == "gateway-key" + assert kwargs["api_base"] == "https://openrouter.ai/api/v1" + + +@pytest.mark.asyncio +async def test_openai_gateway_credentials_reach_json_provider(monkeypatch): + overrides = { + "OPENAI.KEY": "gateway-key", + "OPENAI.API_BASE": "https://gateway.example/v1", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "xiaomi_mimo/mimo-v2-flash") + + assert kwargs["api_key"] == "gateway-key" + assert kwargs["api_base"] == "https://gateway.example/v1" + + +@pytest.mark.asyncio +async def test_json_provider_environment_base_is_request_local(monkeypatch): + monkeypatch.setenv("PUBLICAI_API_BASE", "https://request.example/v1") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) + handler = LiteLLMAIHandler() + monkeypatch.setenv("PUBLICAI_API_BASE", "https://another-request.example/v1") + + kwargs = await _call(handler, "publicai/model") + + assert kwargs["api_base"] == "https://request.example/v1" + + +@pytest.mark.asyncio +async def test_json_provider_blocks_unrelated_openai_environment(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-openai-key") + + kwargs = await _call(LiteLLMAIHandler(), "xiaomi_mimo/mimo-v2-flash") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.asyncio +async def test_native_compatible_provider_environment_ignores_openai_gateway(monkeypatch): + monkeypatch.setenv("DEEPSEEK_API_KEY", "native-deepseek-key") + overrides = { + "OPENAI.KEY": "gateway-key", + "OPENAI.API_BASE": "https://gateway.example/v1", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "deepseek/deepseek-chat") + + assert kwargs["api_key"] == "native-deepseek-key" + assert "api_base" not in kwargs + + +@pytest.mark.asyncio +async def test_native_compatible_provider_setting_ignores_openai_gateway(monkeypatch): + overrides = { + "DEEPSEEK.KEY": "configured-deepseek-key", + "OPENAI.KEY": "gateway-key", + "OPENAI.API_BASE": "https://gateway.example/v1", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "deepseek/deepseek-chat") + + assert kwargs["api_key"] == "configured-deepseek-key" + assert "api_base" not in kwargs + + +@pytest.mark.asyncio +async def test_openai_gateway_base_does_not_override_compatible_provider_environment(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "native-openai-key") + monkeypatch.setenv("DEEPSEEK_API_KEY", "native-deepseek-key") + overrides = {"OPENAI.API_BASE": "https://gateway.example/v1"} + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "deepseek/deepseek-chat") + + assert kwargs["api_key"] == "native-deepseek-key" + assert "api_base" not in kwargs + + +@pytest.mark.parametrize("openrouter_api_base", (None, "")) +@pytest.mark.asyncio +async def test_empty_openrouter_api_base_uses_openrouter_default(monkeypatch, openrouter_api_base): + overrides = { + "OPENROUTER.KEY": "openrouter-key", + "OPENROUTER.API_BASE": openrouter_api_base, + "OPENAI.API_BASE": "https://gateway.example/v1", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "openrouter/model") + + assert kwargs["api_key"] == "openrouter-key" + assert kwargs["api_base"] == "https://openrouter.ai/api/v1" + + +@pytest.mark.asyncio +async def test_native_openrouter_key_uses_openrouter_default(monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "native-openrouter-key") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), "openrouter/model") + + assert kwargs["api_key"] == "native-openrouter-key" + assert kwargs["api_base"] == "https://openrouter.ai/api/v1" + + +@pytest.mark.asyncio +async def test_openrouter_without_key_ignores_openai_gateway(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), "openrouter/model") + + assert "api_base" not in kwargs + + +@pytest.mark.parametrize("model", ("anthropic/claude-x", "openai_like/model", "openrouter/model")) +@pytest.mark.asyncio +async def test_azure_endpoint_is_not_used_for_non_azure_fallback(monkeypatch, model): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.KEY": "azure-key", + "OPENAI.API_BASE": "https://azure.example/v1", + "OPENAI.DEPLOYMENT_ID": "azure-deployment", + "ANTHROPIC.KEY": "anthropic-key", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert "api_base" not in kwargs + assert "deployment_id" not in kwargs + if model.startswith("anthropic/"): + assert kwargs["api_key"] == "anthropic-key" + else: + assert "api_key" not in kwargs + + +@pytest.mark.parametrize("model", ("gpt-4o", "openai/gpt-4o")) +@pytest.mark.asyncio +async def test_deployment_id_follows_the_current_fallback_attempt(monkeypatch, model): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.DEPLOYMENT_ID": "primary-deployment", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + handler = LiteLLMAIHandler() + overrides["OPENAI.DEPLOYMENT_ID"] = "fallback-deployment" + + kwargs = await _call(handler, model) + + assert kwargs["model"] == "azure/gpt-4o" + assert kwargs["deployment_id"] == "fallback-deployment" + + +@pytest.mark.parametrize("deployment_id", ("fallback-deployment", "")) +@pytest.mark.asyncio +async def test_empty_deployment_id_is_not_forwarded(monkeypatch, deployment_id): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.DEPLOYMENT_ID": "primary-deployment", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + handler = LiteLLMAIHandler() + overrides["OPENAI.DEPLOYMENT_ID"] = deployment_id + + kwargs = await _call(handler, "gpt-4o") + + if deployment_id: + assert kwargs["deployment_id"] == deployment_id + else: + assert "deployment_id" not in kwargs + + +@pytest.mark.asyncio +async def test_azure_deployment_id_is_not_used_for_non_azure_probe(monkeypatch): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.DEPLOYMENT_ID": "azure-deployment", + "ANTHROPIC.KEY": "anthropic-key", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + handler = LiteLLMAIHandler() + completion = AsyncMock(return_value=_mock_response()) + + await handler.probe_completion("anthropic/claude-x", _completion=completion) + + assert completion.call_args.kwargs["model"] == "anthropic/claude-x" + assert "deployment_id" not in completion.call_args.kwargs + + +@pytest.mark.parametrize(("setting_path", "model", "expected_key"), ( + ("ANTHROPIC.KEY", "anthropic/claude-sonnet-4-5", "anthropic-key"), + ("COHERE.KEY", "cohere/command-r", "cohere-key"), + ("GROQ.KEY", "groq/llama-3.3-70b-versatile", "groq-key"), + ("SAMBANOVA.KEY", "sambanova/Meta-Llama-3.3-70B-Instruct", "sambanova-key"), + ("REPLICATE.KEY", "replicate/meta/model", "replicate-key"), + ("XAI.KEY", "xai/grok-4", "xai-key"), + ("GOOGLE_AI_STUDIO.GEMINI_API_KEY", "gemini/gemini-2.5-pro", "gemini-key"), + ("DEEPSEEK.KEY", "deepseek/deepseek-chat", "deepseek-key"), + ("ZAI.KEY", "zai/glm-4.5", "zai-key"), + ("DASHSCOPE.KEY", "dashscope/qwen3.8-max", "dashscope-key"), + ("XIAOMI_MIMO.KEY", "xiaomi_mimo/mimo-v2-flash", "xiaomi-key"), + ("DEEPINFRA.KEY", "deepinfra/meta-llama/model", "deepinfra-key"), + ("MISTRAL.KEY", "mistral/mistral-large-latest", "mistral-key"), + ("CODESTRAL.KEY", "codestral/codestral-latest", "codestral-key"), +)) +@pytest.mark.asyncio +async def test_provider_key_is_forwarded_only_for_matching_model(monkeypatch, setting_path, model, expected_key): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({setting_path: expected_key}), + ) + handler = LiteLLMAIHandler() + + matching_kwargs = await _call(handler, model) + openai_kwargs = await _call(handler, "gpt-4o") + + assert matching_kwargs["api_key"] == expected_key + assert openai_kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + +@pytest.mark.asyncio +async def test_provider_specific_endpoints_do_not_cross_models(monkeypatch): + overrides = { + "MOONSHOT.KEY": "moonshot-key", + "MOONSHOT.API_BASE": "https://api.moonshot.cn/v1", + "OLLAMA.API_KEY": "ollama-key", + "OLLAMA.API_BASE": "http://ollama-a:11434", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + handler = LiteLLMAIHandler() + + moonshot_kwargs = await _call(handler, "moonshot/kimi-k3") + ollama_kwargs = await _call(handler, "ollama/llama3") + + assert moonshot_kwargs["api_key"] == "moonshot-key" + assert moonshot_kwargs["api_base"] == "https://api.moonshot.cn/v1" + assert ollama_kwargs["api_key"] == "ollama-key" + assert ollama_kwargs["api_base"] == "http://ollama-a:11434" + + +@pytest.mark.asyncio +async def test_sequential_handlers_keep_request_credentials_isolated(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"MOONSHOT.KEY": "tenant-a", "MOONSHOT.API_BASE": "https://a.example/v1"}), + ) + tenant_a = LiteLLMAIHandler() + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"MOONSHOT.KEY": "tenant-b", "MOONSHOT.API_BASE": "https://b.example/v1"}), + ) + tenant_b = LiteLLMAIHandler() + + tenant_a_kwargs = await _call(tenant_a, "moonshot/kimi-k3") + tenant_b_kwargs = await _call(tenant_b, "moonshot/kimi-k3") + + assert (tenant_a_kwargs["api_key"], tenant_a_kwargs["api_base"]) == ("tenant-a", "https://a.example/v1") + assert (tenant_b_kwargs["api_key"], tenant_b_kwargs["api_base"]) == ("tenant-b", "https://b.example/v1") + + +@pytest.mark.asyncio +async def test_concurrent_handlers_keep_request_credentials_isolated(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENROUTER.KEY": "tenant-a", "OPENROUTER.API_BASE": "https://a.example/v1"}), + ) + tenant_a = LiteLLMAIHandler() + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENROUTER.KEY": "tenant-b", "OPENROUTER.API_BASE": "https://b.example/v1"}), + ) + tenant_b = LiteLLMAIHandler() + calls = [] + + async def capture_call(**kwargs): + await asyncio.sleep(0) + calls.append(kwargs) + return _mock_response() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=capture_call): + await asyncio.gather( + tenant_a.chat_completion(model="openrouter/openai/gpt-4o", system="sys", user="usr"), + tenant_b.chat_completion(model="openrouter/openai/gpt-4o", system="sys", user="usr"), + ) + + credentials = {(call["api_key"], call["api_base"]) for call in calls} + assert credentials == {("tenant-a", "https://a.example/v1"), ("tenant-b", "https://b.example/v1")} + + +@pytest.mark.asyncio +async def test_vertex_routing_is_request_local(monkeypatch): + overrides = { + "VERTEXAI.VERTEX_PROJECT": "request-project", + "VERTEXAI.VERTEX_LOCATION": "asia-east1", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + monkeypatch.setattr(litellm, "api_key", "another-request-key") + + kwargs = await _call(LiteLLMAIHandler(), "vertex_ai/gemini-2.5-pro") + + assert kwargs["vertex_project"] == "request-project" + assert kwargs["vertex_location"] == "asia-east1" + assert "api_key" not in kwargs + + +@pytest.mark.parametrize(("model", "setting_path", "expected_key"), ( + ("anthropic_text/claude-2", "ANTHROPIC.KEY", "anthropic-key"), + ("azure_text/gpt-35-turbo-instruct", "OPENAI.KEY", "azure-key"), + ("ollama_chat/llama3", "OLLAMA.API_KEY", "ollama-key"), + ("text-completion-openai/gpt-3.5-turbo-instruct", "OPENAI.KEY", "openai-key"), + ("vertex_ai_beta/gemini-2.5-pro", "VERTEXAI.VERTEX_PROJECT", "vertex-project"), +)) +@pytest.mark.asyncio +async def test_provider_alias_uses_canonical_request_settings(monkeypatch, model, setting_path, expected_key): + overrides = {setting_path: expected_key} + if model.startswith("azure_text/"): + overrides["OPENAI.API_TYPE"] = "azure" + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), model) + + parameter = "vertex_project" if setting_path == "VERTEXAI.VERTEX_PROJECT" else "api_key" + assert kwargs[parameter] == expected_key + + +@pytest.mark.asyncio +async def test_explicit_azure_model_uses_native_key_outside_azure_mode(monkeypatch): + monkeypatch.setenv("AZURE_API_KEY", "native-azure-key") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.KEY": "openai-key"}), + ) + + kwargs = await _call(LiteLLMAIHandler(), "azure/gpt-4o") + + assert kwargs["api_key"] == "native-azure-key" + + +@pytest.mark.asyncio +async def test_bare_model_provider_alias_uses_canonical_request_settings(monkeypatch): + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({"ANTHROPIC.KEY": "anthropic-key"})) + monkeypatch.setattr(litellm, "get_llm_provider", lambda model: (model, "anthropic_text", None, None)) + + kwargs = await _call(LiteLLMAIHandler(), "claude-2") + + assert kwargs["api_key"] == "anthropic-key" + + +@pytest.mark.asyncio +async def test_bare_model_provider_resolution_is_cached_per_handler(monkeypatch): + resolve_provider = MagicMock(return_value=("gpt-4o", "openai", None, None)) + monkeypatch.setattr(litellm, "get_llm_provider", resolve_provider) + handler = LiteLLMAIHandler() + + await _call(handler, "gpt-4o") + await _call(handler, "gpt-4o") + + resolve_provider.assert_called_once_with(model="gpt-4o") + + +@pytest.mark.asyncio +async def test_bare_model_provider_resolution_only_falls_back_for_bad_requests(monkeypatch): + bad_request = litellm.BadRequestError( + message="provider not found", + model="custom-model", + llm_provider="", + ) + resolve_provider = MagicMock(side_effect=bad_request) + monkeypatch.setattr(litellm, "get_llm_provider", resolve_provider) + + kwargs = await _call(LiteLLMAIHandler(), "custom-model") + + assert kwargs["api_key"] == DUMMY_LITELLM_API_KEY + + resolve_provider.side_effect = RuntimeError("provider resolution failed") + with pytest.raises(RuntimeError, match="provider resolution failed"): + await _call(LiteLLMAIHandler(), "another-custom-model") + + +@pytest.mark.asyncio +async def test_azure_ad_token_and_endpoint_are_request_local(monkeypatch): + overrides = { + "AZURE_AD.CLIENT_ID": "client-id", + "AZURE_AD.API_BASE": "https://azure.example", + "OPENAI.DEPLOYMENT_ID": "azure-deployment", + "OPENAI.API_VERSION": "2026-01-01", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + credential = object() + credentials = [] + tokens = iter(("azure-ad-token-1", "azure-ad-token-2")) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_credential", lambda settings: credential) + + def get_token(received_credential): + credentials.append(received_credential) + return next(tokens) + + monkeypatch.setattr(litellm_handler, "_get_azure_ad_token", get_token) + handler = LiteLLMAIHandler() + + first_kwargs = await _call(handler, "gpt-4o") + second_kwargs = await _call(handler, "gpt-4o") + + assert first_kwargs["model"] == "azure/gpt-4o" + assert first_kwargs["azure_ad_token"] == "azure-ad-token-1" + assert second_kwargs["azure_ad_token"] == "azure-ad-token-2" + assert first_kwargs["api_base"] == "https://azure.example" + assert first_kwargs["api_version"] == "2026-01-01" + assert first_kwargs["deployment_id"] == "azure-deployment" + assert credentials == [credential, credential] + + +def test_azure_ad_token_requires_request_local_credential(): + with pytest.raises(ValueError, match="credential is required"): + litellm_handler._get_azure_ad_token(None) + + +def test_azure_ad_token_error_redacts_provider_details(monkeypatch): + credential = MagicMock() + credential.get_token.side_effect = RuntimeError("provider-secret") + logger = MagicMock() + monkeypatch.setattr(litellm_helpers, "get_logger", lambda: logger) + + with pytest.raises(RuntimeError, match="provider-secret"): + litellm_helpers._get_azure_ad_token(credential) + + logger.error.assert_called_once_with("Failed to get Azure AD token: RuntimeError") + assert "provider-secret" not in logger.error.call_args.args[0] + + +@pytest.mark.asyncio +async def test_health_probe_routes_azure_ad_token_per_request(monkeypatch): + overrides = { + "AZURE_AD.CLIENT_ID": "client-id", + "AZURE_AD.API_BASE": "https://azure.example", + "OPENAI.DEPLOYMENT_ID": "azure-deployment", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + credential = object() + get_token = MagicMock(return_value="probe-token") + monkeypatch.setattr(litellm_handler, "_get_azure_ad_credential", lambda settings: credential) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_token", get_token) + handler = LiteLLMAIHandler() + completion = AsyncMock(return_value=_mock_response()) + + await handler.probe_completion("gpt-4o", _completion=completion) + + completion.assert_awaited_once() + assert completion.call_args.kwargs["model"] == "azure/gpt-4o" + assert completion.call_args.kwargs["azure_ad_token"] == "probe-token" + assert completion.call_args.kwargs["api_base"] == "https://azure.example" + assert completion.call_args.kwargs["deployment_id"] == "azure-deployment" + assert get_token.call_count == 1 + assert all(token_call.args == (credential,) for token_call in get_token.call_args_list) + + +def test_azure_ad_credential_creation_failure_is_logged(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"AZURE_AD.CLIENT_ID": "client-id"}), + ) + logger = MagicMock() + monkeypatch.setattr(litellm_handler, "get_logger", lambda: logger) + secret = "azure-client-secret" + + def fail_to_create_credential(settings): + raise RuntimeError(f"credential setup failed with {secret}") + + monkeypatch.setattr(litellm_handler, "_get_azure_ad_credential", fail_to_create_credential) + + with pytest.raises(RuntimeError, match="credential setup failed"): + LiteLLMAIHandler() + + logger.error.assert_called_once_with("Failed to create Azure AD credential: RuntimeError") + assert secret not in logger.error.call_args.args[0] + + +@pytest.mark.asyncio +async def test_azure_ad_refresh_failure_is_not_retried(monkeypatch): + overrides = { + "AZURE_AD.CLIENT_ID": "client-id", + "AZURE_AD.API_BASE": "https://azure.example", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + credential = object() + get_token = MagicMock(side_effect=RuntimeError("refresh failed")) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_credential", lambda settings: credential) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_token", get_token) + + async def run_inline(function, *args): + return function(*args) + + monkeypatch.setattr(litellm_handler.asyncio, "to_thread", run_inline) + completion = AsyncMock() + handler = LiteLLMAIHandler() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", completion): + with pytest.raises(RuntimeError, match="refresh failed"): + await handler.chat_completion(model="gpt-4o", system="sys", user="usr") + + assert get_token.call_count == 1 + completion.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_azure_ad_uses_openai_endpoint_when_its_endpoint_is_unset(monkeypatch): + overrides = { + "AZURE_AD.CLIENT_ID": "client-id", + "OPENAI.API_BASE": "https://azure.example", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + credential = object() + monkeypatch.setattr(litellm_handler, "_get_azure_ad_credential", lambda settings: credential) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_token", lambda received: "azure-ad-token") + + kwargs = await _call(LiteLLMAIHandler(), "gpt-4o") + + assert kwargs["api_base"] == "https://azure.example" + + +@pytest.mark.asyncio +async def test_azure_ad_preserves_native_non_azure_provider_routing(monkeypatch): + overrides = { + "AZURE_AD.CLIENT_ID": "client-id", + "AZURE_AD.API_BASE": "https://azure.example", + "OPENAI.API_BASE": "https://gateway.example/v1", + "DEEPSEEK.KEY": "deepseek-key", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_credential", lambda settings: object()) + + kwargs = await _call(LiteLLMAIHandler(), "deepseek/deepseek-chat") + + assert kwargs["api_key"] == "deepseek-key" + assert "api_base" not in kwargs + + +@pytest.mark.asyncio +async def test_azure_mode_preserves_explicit_non_openai_provider(monkeypatch): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.KEY": "azure-key", + "ANTHROPIC.KEY": "anthropic-key", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "anthropic/claude-sonnet-4-5") + + assert kwargs["model"] == "anthropic/claude-sonnet-4-5" + assert kwargs["api_key"] == "anthropic-key" + + +@pytest.mark.parametrize(("model", "expected_model"), ( + ("text-completion-openai/gpt-3.5-turbo-instruct", "azure_text/gpt-3.5-turbo-instruct"), + ("azure_text/gpt-3.5-turbo-instruct", "azure_text/gpt-3.5-turbo-instruct"), + ("aiohttp_openai/gpt-4o", "azure/gpt-4o"), +)) +@pytest.mark.asyncio +async def test_azure_mode_routes_openai_aliases_to_azure(monkeypatch, model, expected_model): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({ + "OPENAI.API_TYPE": "azure", + "OPENAI.KEY": "azure-key", + "OPENAI.API_BASE": "https://azure.example", + "OPENAI.API_VERSION": "2026-01-01", + }), + ) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["model"] == expected_model + assert kwargs["api_key"] == "azure-key" + assert kwargs["api_base"] == "https://azure.example" + assert kwargs["api_version"] == "2026-01-01" + + +@pytest.mark.parametrize("model", ( + "text-completion-openai/gpt-3.5-turbo-instruct", + "azure_text/gpt-3.5-turbo-instruct", +)) +@pytest.mark.asyncio +async def test_azure_text_alias_uses_deployment_id_as_model(monkeypatch, model): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.KEY": "azure-key", + "OPENAI.API_BASE": "https://azure.example", + "OPENAI.API_VERSION": "2026-01-01", + "OPENAI.DEPLOYMENT_ID": "azure-text-deployment", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), model) + + assert kwargs["model"] == "azure_text/azure-text-deployment" + assert "deployment_id" not in kwargs + assert kwargs["api_key"] == "azure-key" + assert kwargs["api_base"] == "https://azure.example" + assert kwargs["api_version"] == "2026-01-01" + + +@pytest.mark.asyncio +async def test_explicit_azure_text_uses_deployment_id_outside_azure_mode(monkeypatch): + monkeypatch.setenv("AZURE_API_KEY", "azure-key") + overrides = { + "OPENAI.API_BASE": "https://azure.example", + "OPENAI.API_VERSION": "2026-01-01", + "OPENAI.DEPLOYMENT_ID": "azure-text-deployment", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + + kwargs = await _call(LiteLLMAIHandler(), "azure_text/gpt-3.5-turbo-instruct") + + assert kwargs["model"] == "azure_text/azure-text-deployment" + assert "deployment_id" not in kwargs + assert kwargs["api_key"] == "azure-key" + assert kwargs["api_base"] == "https://azure.example" + assert kwargs["api_version"] == "2026-01-01" + + +@pytest.mark.asyncio +async def test_deployment_id_is_snapshotted_per_request(monkeypatch): + monkeypatch.setenv("AZURE_API_KEY", "azure-key") + overrides = { + "OPENAI.API_BASE": "https://azure.example", + "OPENAI.API_VERSION": "2026-01-01", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + handler = LiteLLMAIHandler() + deployment_reads = 0 + + def changing_deployment_id(self): + nonlocal deployment_reads + deployment_reads += 1 + return f"deployment-{deployment_reads}" + + monkeypatch.setattr(LiteLLMAIHandler, "deployment_id", property(changing_deployment_id)) + + chat_kwargs = await _call(handler, "azure_text/gpt-3.5-turbo-instruct") + completion = AsyncMock(return_value=_mock_response()) + await handler.probe_completion( + "azure_text/gpt-3.5-turbo-instruct", + _completion=completion, + ) + + assert chat_kwargs["model"] == "azure_text/deployment-1" + assert completion.call_args.kwargs["model"] == "azure_text/deployment-2" + assert deployment_reads == 2 + + +@pytest.mark.asyncio +async def test_deployment_id_snapshot_is_reused_for_same_model_retry(monkeypatch): + monkeypatch.setenv("AZURE_API_KEY", "azure-key") + overrides = { + "OPENAI.API_BASE": "https://azure.example", + "OPENAI.API_VERSION": "2026-01-01", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + handler = LiteLLMAIHandler() + deployment_reads = 0 + + def changing_deployment_id(self): + nonlocal deployment_reads + deployment_reads += 1 + return f"deployment-{deployment_reads}" + + monkeypatch.setattr(LiteLLMAIHandler, "deployment_id", property(changing_deployment_id)) + completion = AsyncMock(side_effect=[ + openai.APIError("retry", request=httpx.Request("POST", "https://azure.example"), body=None), + _mock_response(), + ]) + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", completion): + await handler.chat_completion( + model="azure_text/gpt-3.5-turbo-instruct", + system="sys", + user="usr", + ) + + assert [call.kwargs["model"] for call in completion.await_args_list] == [ + "azure_text/deployment-1", + "azure_text/deployment-1", + ] + assert deployment_reads == 1 + + +@pytest.mark.parametrize("model", ( + "text-completion-openai/gpt-3.5-turbo-instruct", + "azure_text/gpt-3.5-turbo-instruct", +)) +@pytest.mark.asyncio +async def test_azure_text_alias_probe_uses_deployment_id_as_model(monkeypatch, model): + overrides = { + "OPENAI.API_TYPE": "azure", + "OPENAI.KEY": "azure-key", + "OPENAI.API_BASE": "https://azure.example", + "OPENAI.API_VERSION": "2026-01-01", + "OPENAI.DEPLOYMENT_ID": "azure-text-deployment", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + completion = AsyncMock(return_value=_mock_response()) + + await LiteLLMAIHandler().probe_completion( + model, + _completion=completion, + ) + + assert completion.call_args.kwargs["model"] == "azure_text/azure-text-deployment" + assert "deployment_id" not in completion.call_args.kwargs + assert completion.call_args.kwargs["api_key"] == "azure-key" + assert completion.call_args.kwargs["api_base"] == "https://azure.example" + assert completion.call_args.kwargs["api_version"] == "2026-01-01" + + +@pytest.mark.parametrize("model", ("gpt-5_thinking", "openai/gpt-5_thinking")) +@pytest.mark.asyncio +async def test_health_probe_normalizes_gpt5_thinking_model(monkeypatch, model): + monkeypatch.setattr(litellm_handler, "get_settings", _make_settings) + completion = AsyncMock(return_value=_mock_response()) + + await LiteLLMAIHandler().probe_completion(model, _completion=completion) + + assert completion.call_args.kwargs["model"] == "openai/gpt-5" + + +@pytest.mark.parametrize("initial_base", (None, "https://handler.example/v1")) +@pytest.mark.asyncio +async def test_image_wait_preserves_native_endpoint_isolation(monkeypatch, initial_base): + from litellm.litellm_core_utils import logging_worker + + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({"GROQ.KEY": "handler-key"})) + if initial_base: + monkeypatch.setenv("GROQ_API_BASE", initial_base) + captured = [] + + def image_wait(*args, **kwargs): + monkeypatch.setenv("GROQ_API_BASE", "https://another-handler.example/v1") + return MagicMock(status_code=200) + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + captured.append(request) + raise TransportReached + + # Run real LiteLLM dispatch/authentication; intercept only image I/O and HTTP transport. + monkeypatch.setattr(litellm_handler.requests, "head", image_wait) + monkeypatch.setattr(httpx.AsyncClient, "send", send) + handler = LiteLLMAIHandler() + try: + if initial_base: + with pytest.raises(TransportReached): + await handler.chat_completion( + model="groq/llama-3.3-70b-versatile", system="sys", user="usr", + img_path="https://image.example/test.png", + ) + assert len(captured) == 1 + assert str(captured[0].url) == f"{initial_base}/chat/completions" + assert captured[0].headers["authorization"] == "Bearer handler-key" + assert json.loads(captured[0].content)["messages"][1]["content"][1] == { + "type": "image_url", "image_url": {"url": "https://image.example/test.png"}, + } + else: + with pytest.raises(ValueError, match="Refusing live api_base environment fallback for provider groq"): + await handler.chat_completion( + model="groq/llama-3.3-70b-versatile", system="sys", user="usr", + img_path="https://image.example/test.png", + ) + assert not captured + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + + +@pytest.mark.asyncio +async def test_azure_ad_token_is_not_resolved_for_explicit_non_azure_provider(monkeypatch): + overrides = { + "AZURE_AD.CLIENT_ID": "client-id", + "ANTHROPIC.KEY": "anthropic-key", + } + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + credential = object() + get_token = MagicMock(side_effect=RuntimeError("Azure AD unavailable")) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_credential", lambda settings: credential) + monkeypatch.setattr(litellm_handler, "_get_azure_ad_token", get_token) + + kwargs = await _call(LiteLLMAIHandler(), "anthropic/claude-sonnet-4-5") + + assert kwargs["model"] == "anthropic/claude-sonnet-4-5" + assert kwargs["api_key"] == "anthropic-key" + get_token.assert_not_called() + + +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("route", ("global", "captured", "host")) +@pytest.mark.parametrize("timing", ("before", "after")) +@pytest.mark.asyncio +async def test_native_gdc_rejects_uncaptured_routing(monkeypatch, entrypoint, route, timing): + from litellm.litellm_core_utils import logging_worker + + endpoint = "https://owned.example/v1/projects/owned/locations/local/chat/completions" + monkeypatch.setenv("GDC_API_KEY", "owned-gdc-token") + if route == "captured": + monkeypatch.setenv("GDC_API_BASE", endpoint) + elif route == "host": + monkeypatch.setenv("GDC_API_BASE", "https://owned.example") + def set_foreign_route(): + monkeypatch.setattr( + litellm, "gdc_api_base", + "https://foreign.example/v1/projects/foreign/locations/other/chat/completions", + ) + monkeypatch.setattr(litellm, "vertex_project", "foreign") + monkeypatch.setattr(litellm, "vertex_location", "other") + + if timing == "before": + set_foreign_route() + handler = LiteLLMAIHandler() + if timing == "after": + set_foreign_route() + if route == "captured": + monkeypatch.setenv("GDC_API_BASE", "https://changed.example") + sent = [] + + async def send(client, request, **kwargs): + sent.append(request) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "gemini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + async def invoke(): + if entrypoint == "chat": + await handler.chat_completion("gdc/gemini", "sys", "usr") + else: + await handler.probe_completion("gdc/gemini") + + if route == "captured": + await invoke() + assert len(sent) == 1 + assert str(sent[0].url) == endpoint + assert sent[0].headers["Authorization"] == "Bearer owned-gdc-token" + assert sent[0].headers["x-goog-user-project"] == "projects/owned" + else: + with pytest.raises(ValueError, match="(?i)gdc|routing|process-wide"): + await invoke() + assert sent == [] + finally: + await asyncio.sleep(0) + try: + await worker.flush() + finally: + await worker.stop() + + +@pytest.mark.parametrize("provider", ("mistral", "deepseek", "groq")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("captured_key", (None, "owned-key", DUMMY_LITELLM_API_KEY)) +@pytest.mark.parametrize(("header_name", "authorization"), ( + ("Authorization", "Bearer owned-header"), ("aUtHoRiZaTiOn", "Basic owned-header"), +)) +@pytest.mark.asyncio +async def test_native_raw_guard_preserves_explicit_authorization( + monkeypatch, provider, entrypoint, captured_key, header_name, authorization, +): + from litellm.litellm_core_utils import logging_worker + + settings = _make_settings({ + f"{provider.upper()}.KEY": captured_key, + "LITELLM.EXTRA_HEADERS": json.dumps({header_name: authorization}), + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "api_key", "foreign-key") + sent = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + sent.append(request) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + with pytest.raises(TransportReached): + if entrypoint == "chat": + await handler.chat_completion(f"{provider}/test-model", "sys", "usr") + else: + await handler.probe_completion(f"{provider}/test-model") + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + assert len(sent) == 1 + if captured_key: + # A real key retains native precedence, including a literal dummy key. + assert f"Bearer {captured_key}" in sent[0].headers.get_list("authorization") + else: + assert sent[0].headers.get_list("authorization") == [authorization] + assert "foreign-key" not in str(sent[0].headers) + assert litellm_handler._raw_api_key_guard_provider.get() is None + + +@pytest.mark.parametrize("control", ("active", "no-header", "changed-key", "provider", "override", "outside")) +@pytest.mark.asyncio +async def test_raw_guard_bridge_boundaries(monkeypatch, control): + from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig + + litellm_handler._install_raw_api_key_guard_bridge() + + class InheritedConfig(OpenAIGPTConfig): + pass + + class OverrideConfig(OpenAIGPTConfig): + def validate_environment(self, *args, **kwargs): + return super().validate_environment(*args, **kwargs) + + config = OverrideConfig() if control == "override" else InheritedConfig() + key = "native-key" if control == "changed-key" else DUMMY_LITELLM_API_KEY + + async def completion(**kwargs): + assert "_raw_api_key_guard" not in kwargs + return config.validate_environment( + headers={} if control == "no-header" else {"authorization": "Basic owned"}, + model="test-model", messages=[], optional_params={}, + litellm_params={"custom_llm_provider": "deepseek" if control == "provider" else "mistral"}, + api_key=key, + ) + + result = await LiteLLMAIHandler()._acompletion( + _completion=completion, _raw_api_key_guard=control != "outside", model="mistral/test-model", + api_key=DUMMY_LITELLM_API_KEY, headers={"authorization": "Basic owned"}, + ) + if control == "active": + assert result == {"authorization": "Basic owned", "Content-Type": "application/json"} + else: + assert result["Authorization"] == f"Bearer {key}" + assert litellm_handler._raw_api_key_guard_provider.get() is None + + +@pytest.mark.asyncio +async def test_raw_guard_nested_and_concurrent_contexts(monkeypatch): + from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig + + handler = LiteLLMAIHandler() + both_started = asyncio.Event() + started = [] + + async def inner(**kwargs): + assert litellm_handler._raw_api_key_guard_provider.get() is None + return OpenAIGPTConfig().validate_environment( + headers={"Authorization": "Basic inner"}, model="test-model", messages=[], optional_params={}, + litellm_params={"custom_llm_provider": "mistral"}, api_key=DUMMY_LITELLM_API_KEY, + ) + + async def outer(**kwargs): + started.append(kwargs["headers"]["Authorization"]) + if len(started) == 2: + both_started.set() + await both_started.wait() + assert litellm_handler._raw_api_key_guard_provider.get() == "mistral" + result = await handler._acompletion( + _completion=inner, model="mistral/test-model", api_key=DUMMY_LITELLM_API_KEY, + ) + assert result["Authorization"] == f"Bearer {DUMMY_LITELLM_API_KEY}" + assert litellm_handler._raw_api_key_guard_provider.get() == "mistral" + return OpenAIGPTConfig().validate_environment( + headers=dict(kwargs["headers"]), model="test-model", messages=[], optional_params={}, + litellm_params={"custom_llm_provider": "mistral"}, api_key=DUMMY_LITELLM_API_KEY, + ) + + results = await asyncio.gather(*( + handler._acompletion( + _completion=outer, _raw_api_key_guard=True, model="mistral/test-model", + api_key=DUMMY_LITELLM_API_KEY, headers={"Authorization": value}, + ) + for value in ("Basic first", "Basic second") + )) + assert [result["Authorization"] for result in results] == ["Basic first", "Basic second"] + assert litellm_handler._raw_api_key_guard_provider.get() is None + + +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("source", ("settings", "environment")) +@pytest.mark.parametrize("complete_url", (False, True)) +@pytest.mark.asyncio +async def test_native_gdc_captured_host_routing(monkeypatch, entrypoint, source, complete_url): + from litellm.litellm_core_utils import logging_worker + + endpoint = "https://owned.example/v1/projects/owned/locations/local/chat/completions" + monkeypatch.setenv("GDC_API_BASE", endpoint if complete_url else "https://owned.example") + monkeypatch.setenv("GDC_API_KEY", "owned-token") + project = "ignored" if complete_url else "owned" + if source == "settings": + settings = _make_settings({ + "VERTEXAI.VERTEX_PROJECT": project, "VERTEXAI.VERTEX_LOCATION": "local", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + else: + monkeypatch.setenv("VERTEXAI_PROJECT", project) + monkeypatch.setenv("VERTEX_LOCATION", "local") + handler = LiteLLMAIHandler() + sent, boundaries = [], [] + loop = asyncio.get_running_loop() + executor = loop.run_in_executor + + def dispatch(pool, function, *args): + if not boundaries: + boundaries.append(function) + monkeypatch.setattr(litellm, "vertex_project", "foreign") + monkeypatch.setattr(litellm, "vertex_location", "other") + monkeypatch.setattr(litellm, "gdc_api_base", "https://foreign.example") + return executor(pool, function, *args) + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + sent.append(request) + raise TransportReached + + monkeypatch.setattr(loop, "run_in_executor", dispatch) + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + with pytest.raises(TransportReached): + if entrypoint == "chat": + await handler.chat_completion("gdc/gemini", "sys", "usr") + else: + await handler.probe_completion("gdc/gemini") + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + assert len(boundaries) == len(sent) == 1 + assert str(sent[0].url) == endpoint + assert sent[0].headers["authorization"] == "Bearer owned-token" + assert sent[0].headers["x-goog-user-project"] == "projects/owned" + assert "vertex_" not in sent[0].content.decode() + + +@pytest.mark.parametrize(("base", "project", "location"), ( + (None, "owned", "local"), ("https://owned.example", None, "local"), + ("https://owned.example", "owned", None), ("https://owned.example", "bad/project", "local"), + ("https://owned.example", "owned", "bad/location"), +)) +@pytest.mark.asyncio +async def test_gdc_invalid_host_routing_fails_before_auth(monkeypatch, base, project, location): + from litellm.llms.gdc.chat.transformation import GDCGeminiConfig + + if base: + monkeypatch.setenv("GDC_API_BASE", base) + monkeypatch.setenv("GDC_API_KEY", json.dumps({"type": "gdch_service_account"})) + settings = _make_settings({"VERTEXAI.VERTEX_PROJECT": project, "VERTEXAI.VERTEX_LOCATION": location}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "gdc_api_base", "https://foreign.example") + monkeypatch.setattr(litellm, "vertex_project", "foreign") + monkeypatch.setattr(litellm, "vertex_location", "other") + auth = MagicMock(side_effect=AssertionError("invalid routing must not resolve credentials")) + monkeypatch.setattr(GDCGeminiConfig, "_load_creds_from_key", auth) + with pytest.raises((ValueError, litellm.AuthenticationError)): + await handler.probe_completion("gdc/gemini") + auth.assert_not_called() + + +@pytest.mark.parametrize("configured", (None, "settings-location")) +@pytest.mark.asyncio +async def test_vertex_location_alias_precedence(monkeypatch, configured): + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + settings = _make_settings({"VERTEXAI.VERTEX_LOCATION": configured}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + monkeypatch.setenv("VERTEXAI_LOCATION", "primary-location") + monkeypatch.setenv("VERTEX_LOCATION", "secondary-location") + handler = LiteLLMAIHandler() + monkeypatch.setenv("VERTEXAI_LOCATION", "foreign-primary") + monkeypatch.setenv("VERTEX_LOCATION", "foreign-secondary") + params = await _call(handler, "vertex_ai/gemini-2.5-pro") + expected = configured or "primary-location" + assert VertexBase.get_vertex_ai_location(dict(params)) == expected + assert VertexBase.safe_get_vertex_ai_location(params) == expected + + +@pytest.mark.asyncio +async def test_gdc_handlers_keep_separate_routing(monkeypatch): + monkeypatch.setenv("GDC_API_BASE", "https://owned.example") + monkeypatch.setenv("GDC_API_KEY", "owned-token") + monkeypatch.setenv("VERTEXAI_PROJECT", "first") + monkeypatch.setenv("VERTEXAI_LOCATION", "local") + first = LiteLLMAIHandler() + monkeypatch.setenv("VERTEXAI_PROJECT", "second") + second = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "vertex_project", "foreign") + for handler, project in ((first, "first"), (second, "second"), (first, "first")): + params = await _call(handler, "gdc/gemini") + assert params["api_base"] == f"https://owned.example/v1/projects/{project}/locations/local/chat/completions" + assert "vertex_credentials" not in params + + +@pytest.mark.parametrize("complete_url", (False, True)) +@pytest.mark.asyncio +async def test_native_gdc_service_account_host_routing(monkeypatch, complete_url): + import requests + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import ec + from google.auth import jwt + from litellm.litellm_core_utils import logging_worker + + private_key = ec.generate_private_key(ec.SECP256R1()).private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption(), + ).decode() + endpoint = "https://owned.example/v1/projects/routing-project/locations/local/chat/completions" + monkeypatch.setenv("GDC_API_BASE", endpoint if complete_url else "https://owned.example") + monkeypatch.setenv("GDC_API_KEY", json.dumps({ + "type": "gdch_service_account", "format_version": "1", "private_key_id": "test-key", + "private_key": private_key, "name": "test-service", "project": "identity-project", + "token_uri": "https://identity.example/authenticate", + })) + settings = _make_settings({ + "VERTEXAI.VERTEX_PROJECT": "routing-project", "VERTEXAI.VERTEX_LOCATION": "local", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + handler = LiteLLMAIHandler() + exchanged, sent = [], [] + + def token_send(session, request, **kwargs): + assert request.url == "https://identity.example/authenticate" + body = json.loads(request.body) + assert body["audience"] == "https://owned.example" + claims = jwt.decode(body["subject_token"], verify=False) + assert claims["iss"] == "system:serviceaccount:identity-project:test-service" + assert claims["aud"] == request.url + exchanged.append(body) + response = requests.Response() + response.status_code = 200 + response._content = json.dumps({"access_token": "exchanged-token", "expires_in": 3600}).encode() + return response + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + sent.append(request) + raise TransportReached + + monkeypatch.setattr(requests.Session, "send", token_send) + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + with pytest.raises(TransportReached): + await handler.probe_completion("gdc/gemini") + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + assert len(exchanged) == len(sent) == 1 + assert str(sent[0].url) == endpoint + assert sent[0].headers["authorization"] == "Bearer exchanged-token" + assert sent[0].headers["x-goog-user-project"] == "projects/routing-project" + + +@pytest.fixture +def header_only_native_auth(monkeypatch): + from azure.core.pipeline.transport import RequestsTransport + + for name in (*litellm_handler.AZURE_OIDC_ENV_VARS, *litellm_handler.AZURE_OIDC_AUTH_ENV_VARS): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + for name in ("callbacks", "success_callback", "failure_callback", "input_callback", + "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, [], raising=False) + + def unexpected_auth_http(*args, **kwargs): + raise AssertionError("Unexpected Azure auth HTTP in header-only test") + + monkeypatch.setattr(RequestsTransport, "send", unexpected_auth_http) + + +@pytest.mark.parametrize("route", ( + "neosantara", "tensormesh", "parasail", "meta", "pinstripes", + "xai", "ragflow", "ragflow_agent", "azure_ai", "azure_ai_key", +)) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize("captured_key", (None, "owned-key", DUMMY_LITELLM_API_KEY)) +@pytest.mark.parametrize(("header_name", "authorization"), ( + ("Authorization", "Bearer owned-header"), ("authorization", "Basic owned-header"), + ("aUtHoRiZaTiOn", "Bearer owned-header"), +)) +@pytest.mark.asyncio +async def test_native_guard_header_only_equivalence( + monkeypatch, header_only_native_auth, route, entrypoint, captured_key, header_name, authorization, +): + from litellm.litellm_core_utils import logging_worker + + if route.startswith("ragflow"): + kind = "agent" if route == "ragflow_agent" else "chat" + model = f"ragflow/{kind}/owned/test-model" + base_variable, base, key_variable = "RAGFLOW_API_BASE", "https://owned.example", "RAGFLOW_API_KEY" + endpoint = f"{base}/api/v1/{kind}s_openai/owned/chat/completions" + elif route.startswith("azure_ai"): + model = "azure_ai/test-model" + base_variable, key_variable = "AZURE_AI_API_BASE", "AZURE_AI_API_KEY" + base = "https://owned.services.ai.azure.com" if route == "azure_ai_key" else "https://owned.example" + endpoint = base + ("/models/chat/completions" if route == "azure_ai_key" else "/chat/completions") + elif route == "xai": + model = "xai/test-model" + base_variable, base, key_variable = "XAI_API_BASE", "https://owned.example/v1", "XAI_API_KEY" + endpoint = base + "/chat/completions" + else: + config = litellm_handler.JSONProviderRegistry.get(route) + model = f"{route}/responses/test-model" + base_variable, base, key_variable = config.api_base_env, "https://owned.example/v1", config.api_key_env + endpoint = base + "/responses" + monkeypatch.setenv(base_variable, base) + if captured_key: + monkeypatch.setenv(key_variable, captured_key) + settings = _make_settings({"LITELLM.EXTRA_HEADERS": json.dumps({header_name: authorization})}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + handler = LiteLLMAIHandler() + snapshot = dict(handler._request_headers) + # Unlike api_key, the base handler did not forward this unrelated key. + monkeypatch.setattr(litellm, "openai_key", "foreign-key") + sent = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + sent.append(request) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + # Keep native dispatch/authentication as the oracle; stop only at HTTP. + with pytest.raises(TransportReached): + await litellm.acompletion( + model=model, messages=[{"role": "user", "content": "usr"}], + extra_headers={header_name: authorization}, + ) + with pytest.raises(TransportReached): + if entrypoint == "chat": + await handler.chat_completion(model, "sys", "usr") + else: + await handler.probe_completion(model) + finally: + try: + await asyncio.sleep(0) + await worker.flush() + finally: + await worker.stop() + assert len(sent) == 2 + for request in sent: + assert str(request.url) == endpoint + assert json.loads(request.content)["model"] == "test-model" + assert "foreign-key" not in str(request.headers) + for name in ("authorization", "api-key"): + assert sent[1].headers.get_list(name) == sent[0].headers.get_list(name) + if captured_key is None: + assert sent[1].headers.get_list("authorization") == [authorization] + assert sent[1].headers.get_list("api-key") == [] + assert handler._request_headers == snapshot + assert litellm_handler._raw_api_key_guard_provider.get() is None + assert litellm_handler._raw_api_key_guard_auth.get() is None + + +def test_json_responses_native_config_contract(): + from litellm.llms.openai_like.dynamic_config import create_responses_config_class + from litellm.utils import ProviderConfigManager + + registry = litellm_handler.JSONProviderRegistry + providers = {name for name in registry.list_providers() if registry.supports_responses_api(name)} + assert providers == {"neosantara", "tensormesh", "parasail", "empiriolabs", "meta", "pinstripes"} + # Pinned native HTTP client selection rejects this registry-only provider + # before authentication; do not claim transport coverage for that path. + assert providers - {provider.value for provider in litellm.LlmProviders} == {"empiriolabs"} + for provider in providers: + # The native manager checks Python overrides before the JSON registry. + selected = ProviderConfigManager.get_provider_responses_api_config(provider, model="test-model") + expected = create_responses_config_class(registry.get(provider)) + assert type(selected) is expected + assert type(selected).validate_environment is expected.validate_environment + + +@pytest.mark.parametrize("control", ( + "active", "initial-generic", "initial-empty", "no-guard", "changed-key", "multiple", + "chat", "non-json", "no-header", +)) +@pytest.mark.asyncio +async def test_json_guard_normalization_boundaries(monkeypatch, header_only_native_auth, control): + if control in ("initial-generic", "initial-empty"): + monkeypatch.setattr(litellm, "api_key", "initial-key" if control == "initial-generic" else "") + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "api_key", "late-key") + headers = {} if control == "no-header" else {"authorization": "Basic owned"} + if control == "multiple": + headers["AUTHORIZATION"] = "Basic duplicate" + initial_headers = dict(headers) + model = "parasail/responses/test-model" + if control == "chat": + model = "parasail/test-model" + elif control == "non-json": + model = "openai/responses/test-model" + + async def completion(**kwargs): + assert "_raw_api_key_guard" not in kwargs + return kwargs["headers"] + + result = await handler._acompletion( + _completion=completion, _raw_api_key_guard=control != "no-guard", model=model, + api_key="real-key" if control == "changed-key" else DUMMY_LITELLM_API_KEY, headers=headers, + ) + assert result == ({"Authorization": "Basic owned"} if control in ("active", "initial-empty") else initial_headers) + assert headers == initial_headers + + +@pytest.mark.parametrize("raise_error", (False, True)) +@pytest.mark.asyncio +async def test_raw_guard_auth_snapshot_nested_concurrency(monkeypatch, header_only_native_auth, raise_error): + first = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "api_key", "initial-key") + second = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "api_key", "late-key") + started = [] + both_started = asyncio.Event() + + class ExpectedError(Exception): + pass + + async def inner(**kwargs): + assert litellm_handler._raw_api_key_guard_auth.get() is None + assert litellm_handler._raw_api_key_guard_provider.get() is None + + async def invoke(handler, expected): + async def completion(**kwargs): + state = litellm_handler._raw_api_key_guard_auth.get() + assert state["generic_key"] is expected + started.append(expected) + if len(started) == 2: + both_started.set() + await both_started.wait() + await handler._acompletion(_completion=inner, model="xai/test-model", api_key=DUMMY_LITELLM_API_KEY) + assert litellm_handler._raw_api_key_guard_auth.get() is state + if raise_error: + raise ExpectedError + return state["generic_key"] + + try: + return await handler._acompletion( + _completion=completion, _raw_api_key_guard=True, model="xai/test-model", + api_key=DUMMY_LITELLM_API_KEY, headers={"Authorization": "Basic owned"}, + ) + except ExpectedError: + return expected + finally: + # Check cleanup inside the owning task, not only its parent. + assert litellm_handler._raw_api_key_guard_auth.get() is None + assert litellm_handler._raw_api_key_guard_provider.get() is None + + assert await asyncio.gather(invoke(first, False), invoke(second, True)) == [False, True] + + +@pytest.mark.parametrize("provider", ("ragflow", "xai", "azure_ai")) +@pytest.mark.parametrize("control", ( + "active", "initial-generic", "changed-key", "no-guard", "provider-mismatch", "subclass", "multiple", +)) +@pytest.mark.asyncio +async def test_raw_override_guard_boundaries(monkeypatch, header_only_native_auth, provider, control): + from functools import wraps + + from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig + from litellm.llms.ragflow.chat.transformation import RAGFlowConfig + from litellm.llms.xai.chat.transformation import XAIChatConfig + + config = {"azure_ai": AzureAIStudioConfig, "ragflow": RAGFlowConfig, "xai": XAIChatConfig}[provider] + native = config.validate_environment + signature = inspect.signature(native) + observed_keys = [] + + @wraps(native) + def tracked(self, *args, **kwargs): + observed_keys.append(signature.bind(self, *args, **kwargs).arguments["api_key"]) + return native(self, *args, **kwargs) + + monkeypatch.setattr(config, "validate_environment", tracked) + if control == "initial-generic": + monkeypatch.setattr(litellm, "api_key", "initial-key") + handler = LiteLLMAIHandler() + monkeypatch.setattr(litellm, "api_key", "late-key") + headers = {"authorization": "Basic owned-header"} + if control == "multiple": + headers["AUTHORIZATION"] = "Basic duplicate" + initial_headers = dict(headers) + model = "ragflow/chat/owned/test-model" if provider == "ragflow" else f"{provider}/test-model" + key = "changed-key" if control == "changed-key" else DUMMY_LITELLM_API_KEY + results = [] + + async def completion(**kwargs): + assert "_raw_api_key_guard" not in kwargs + selected = type("DerivedConfig", (config,), {}) if control == "subclass" else config + params = {"custom_llm_provider": "other" if control == "provider-mismatch" else provider} + results.append(selected().validate_environment( + headers=kwargs["headers"], model=model, messages=[], optional_params={}, + litellm_params=params, api_key=kwargs["api_key"], api_base="https://owned.example", + )) + + await handler._acompletion( + _completion=completion, _raw_api_key_guard=control != "no-guard", + model=model, api_key=key, headers=headers, + ) + assert observed_keys == [key] + if control == "active": + assert httpx.Headers(results[0]).get_list("authorization") == ["Basic owned-header"] + assert headers == initial_headers + else: + assert f"Bearer {key}" in httpx.Headers(results[0]).get_list("authorization") + litellm_handler._install_raw_api_key_guard_override_bridge(provider) + installed = config.validate_environment + litellm_handler._install_raw_api_key_guard_override_bridge(provider) + assert config.validate_environment is installed + assert litellm_handler._raw_api_key_guard_provider.get() is None + assert litellm_handler._raw_api_key_guard_auth.get() is None + + +@pytest.mark.parametrize(("provider", "environment", "params", "initial_global", "refresh", "expected"), ( + ("ragflow", {}, {}, None, False, True), + ("xai", {}, {}, None, False, True), + ("xai", {}, {"use_xai_oauth": True}, None, False, False), + ("xai", {}, {"use_xai_oauth": "false"}, None, False, False), + ("xai", {}, {}, "xai_key", False, False), + ("azure_ai", {}, {}, None, False, True), + ("azure_ai", {}, {}, "azure_key", False, False), + ("azure_ai", {}, {}, None, True, False), + ("azure_ai", {"AZURE_AD_TOKEN": "initial-token"}, {}, None, False, False), + ("azure_ai", {"AZURE_OPENAI_AD_TOKEN": "sdk-only-alias"}, {}, None, False, True), + ("azure_ai", {"AZURE_API_KEY": "initial-key"}, {}, None, False, False), + ("azure_ai", {"AZURE_OPENAI_API_KEY": "initial-key"}, {}, None, False, False), + ("azure_ai", {}, {"azure_ad_token": "request-token"}, None, False, False), + ("azure_ai", {}, {"azure_ad_token_provider": object()}, None, False, False), + ("azure_ai", {"AZURE_TENANT_ID": "tenant"}, {}, None, False, True), + ("azure_ai", {"AZURE_CLIENT_ID": "client"}, {}, None, False, True), + ("azure_ai", {"AZURE_TENANT_ID": "tenant", "AZURE_CLIENT_ID": "client"}, + {"client_secret": "secret"}, None, False, False), + ("azure_ai", {"AZURE_CLIENT_SECRET": "secret"}, + {"tenant_id": "tenant", "client_id": "client"}, None, False, False), + ("azure_ai", {"AZURE_USERNAME": "user", "AZURE_PASSWORD": "password"}, + {"client_id": "client"}, None, False, False), + ("azure_ai", {"AZURE_CLIENT_ID": "client"}, + {"azure_username": "user", "azure_password": "password"}, None, False, False), +)) +def test_raw_guard_initial_auth_selection( + monkeypatch, header_only_native_auth, provider, environment, params, initial_global, refresh, expected, +): + for name, value in environment.items(): + monkeypatch.setenv(name, value) + if initial_global: + monkeypatch.setattr(litellm, initial_global, "initial-key") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", refresh) + handler = LiteLLMAIHandler() + for name in environment: + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", not refresh) + token = litellm_handler._raw_api_key_guard_auth.set(handler._raw_guard_auth_snapshot) + try: + assert bool(litellm_handler._raw_guard_has_header_only_auth( + provider, params, {"Authorization": "Basic owned"}, + )) is expected + if provider == "azure_ai": + for name in ("api-key", "API-KEY"): + assert not litellm_handler._raw_guard_has_header_only_auth( + provider, params, {"Authorization": "Basic owned", name: ""}, + ) + finally: + litellm_handler._raw_api_key_guard_auth.reset(token) + + +@pytest.fixture +async def native_endpoint_runtime(monkeypatch): + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.litellm_core_utils import logging_worker + + cache = LLMClientCache() + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", cache) + monkeypatch.setenv("LITELLM_RUST", "false") + for name in ("callbacks", "input_callback", "success_callback", "failure_callback", + "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, []) + + def deny_sync(*args, **kwargs): + raise AssertionError("Unexpected synchronous HTTP") + + monkeypatch.setattr(httpx.Client, "send", deny_sync) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + yield + finally: + await asyncio.sleep(0) + try: + await worker.flush() + finally: + await worker.stop() + for client in {id(value): value for value in cache.cache_dict.values()}.values(): + result = client.close() + if inspect.isawaitable(result): + await result + + +@pytest.mark.usefixtures("native_endpoint_runtime") +@pytest.mark.parametrize(("model", "ignored", "effective", "default_url"), ( + ("mistral/mistral-small-latest", "MISTRAL_API_BASE", "MISTRAL_AZURE_API_BASE", + "https://api.mistral.ai/v1/chat/completions"), + ("baseten/Qwen3-Coder", "BASETEN_API_BASE", None, + "https://inference.baseten.co/v1/chat/completions"), + ("baseten/abc12345", "BASETEN_API_BASE", None, + "https://model-abc12345.api.baseten.co/environments/production/sync/v1/chat/completions"), + ("volcengine/doubao-pro-32k", "ARK_API_BASE", "VOLCENGINE_API_BASE", + "https://ark.cn-beijing.volces.com/api/v3/chat/completions"), + ("volcengine/responses/doubao-pro-32k", "ARK_API_BASE", "VOLCENGINE_API_BASE", + "https://ark.cn-beijing.volces.com/api/v3/responses"), +)) +@pytest.mark.parametrize("selection", ("absent", "ignored", "effective", "both", "ignored_late")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.asyncio +async def test_native_initial_endpoint(monkeypatch, model, ignored, effective, default_url, selection, entrypoint): + if effective is None and selection in ("effective", "both"): + pytest.skip("No separate effective environment endpoint for this native provider") + provider = model.split("/", 1)[0] + monkeypatch.setenv(provider.upper() + "_API_KEY", "owned-key") + if selection in ("ignored", "both"): + monkeypatch.setenv(ignored, "https://ignored.example/v1") + if selection in ("effective", "both"): + monkeypatch.setenv(effective, "https://effective.example/v1") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings()) + seen = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + seen.append(request) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + with pytest.raises(TransportReached): + await litellm.acompletion(model=model, messages=[{"role": "user", "content": "usr"}]) + suffix = "/api/v3/responses" if default_url.endswith("/responses") else "/chat/completions" + expected = "https://effective.example/v1" + suffix if selection in ("effective", "both") else default_url + assert str(seen[0].url) == expected + assert seen[0].headers["Authorization"] == "Bearer owned-key" + handler = LiteLLMAIHandler() + if selection == "ignored_late": + monkeypatch.setenv(ignored, "https://ignored.example/v1") + with pytest.raises(TransportReached): + if entrypoint == "chat": + await handler.chat_completion(model, "sys", "usr") + else: + await handler.probe_completion(model) + assert len(seen) == 2 + assert str(seen[1].url) == str(seen[0].url) + for header in ("Authorization", "api-key"): + assert seen[1].headers.get_list(header) == seen[0].headers.get_list(header) + + +@pytest.mark.usefixtures("native_endpoint_runtime") +@pytest.mark.parametrize("selection", ("absent", "ignored", "effective", "both", "ignored_late")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.asyncio +async def test_native_vertex_initial_endpoint(monkeypatch, tmp_path, selection, entrypoint): + import importlib + + import google.auth + from google.auth.transport.requests import Request + from google.oauth2.credentials import Credentials + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM + + credentials = Credentials("owned-token") + monkeypatch.setattr(credentials, "refresh", lambda request: None) + monkeypatch.setenv("CLOUDSDK_CONFIG", str(tmp_path)) + monkeypatch.setattr(google.auth, "default", lambda *args, **kwargs: (credentials, "owned-project")) + monkeypatch.setattr(litellm_handler, "_load_vertex_default_adc", lambda snapshot, project: (credentials, project)) + monkeypatch.setattr(importlib.import_module("litellm.main"), "vertex_chat_completion", VertexLLM()) + + class UnexpectedAuthHTTP(BaseException): + pass + + def deny_auth(*args, **kwargs): + raise UnexpectedAuthHTTP("Unexpected Google authentication HTTP") + + monkeypatch.setattr(Request, "__call__", deny_auth) + monkeypatch.setenv("VERTEXAI_PROJECT", "owned-project") + monkeypatch.setenv("VERTEXAI_LOCATION", "us-central1") + if selection in ("ignored", "both"): + monkeypatch.setenv("VERTEX_API_BASE", "https://ignored.example") + if selection in ("effective", "both"): + monkeypatch.setenv("VERTEXAI_API_BASE", "https://effective.example") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings()) + seen = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + seen.append(request) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + model = "vertex_ai/gemini-2.5-flash" + with pytest.raises(TransportReached): + await litellm.acompletion(model=model, messages=[{"role": "user", "content": "usr"}]) + host = "effective.example" if selection in ("effective", "both") else "us-central1-aiplatform.googleapis.com" + assert seen[0].url.host == host + assert seen[0].url.path == ( + "/v1/projects/owned-project/locations/us-central1/publishers/google/models/gemini-2.5-flash:generateContent" + ) + assert seen[0].headers["Authorization"] == "Bearer owned-token" + handler = LiteLLMAIHandler() + if selection == "ignored_late": + monkeypatch.setenv("VERTEX_API_BASE", "https://ignored.example") + with pytest.raises(TransportReached): + if entrypoint == "chat": + await handler.chat_completion(model, "sys", "usr") + else: + await handler.probe_completion(model) + assert len(seen) == 2 + assert seen[1].url == seen[0].url + assert seen[1].headers.get_list("Authorization") == seen[0].headers.get_list("Authorization") + + +@pytest.mark.parametrize(("route", "resource"), ( + ("", "foundation-model/anthropic.claude-3-haiku-20240307-v1:0"), + ("invoke/", "foundation-model/anthropic.claude-3-haiku-20240307-v1:0"), +) + tuple( + (prefix + spec + "/", "imported-model/example") + for prefix in ("", "invoke/") + for spec in ("llama", "deepseek_r1", "openai", "qwen2", "qwen3", "moonshot", "nova-2", "nova") +)) +@pytest.mark.usefixtures("native_endpoint_runtime") +@pytest.mark.asyncio +async def test_arn_routing_region(monkeypatch, route, resource): + for name in ("AWS_REGION_NAME", "AWS_REGION", "AWS_DEFAULT_REGION", "DEFAULT_REGION"): + monkeypatch.delenv(name, raising=False) + credentials = { + "aws_access_key_id": "owned-key", "aws_secret_access_key": "owned-secret", + "aws_session_token": "owned-session", + } + for key, value in credentials.items(): + monkeypatch.setenv(key.upper(), value) + monkeypatch.setenv("AWS_USE_IMDS", "false") + settings = _make_settings() + # Imported OpenAI reasoning models reject the default temperature parameter. + settings.config.custom_reasoning_model = route == "invoke/openai/" + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + seen = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + seen.append(request) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + handler = LiteLLMAIHandler() + original_credentials = dict(handler._aws_active_creds) + for region in ("eu-west-1", "us-east-1"): + account = "" if resource.startswith("foundation-model/") else "123456789012" + model = f"bedrock/{route}arn:aws:bedrock:{region}:{account}:{resource}" + with pytest.raises(TransportReached): + await litellm.acompletion(model=model, messages=[{"role": "user", "content": "usr"}], **credentials) + native = seen[-1] + assert native.url.host == f"bedrock-runtime.{region}.amazonaws.com" + assert native.headers["Authorization"].startswith("AWS4-HMAC-SHA256 ") + assert "Credential=owned-key/" in native.headers["Authorization"] + assert f"/{region}/bedrock/aws4_request" in native.headers["Authorization"] + assert native.headers["X-Amz-Security-Token"] == "owned-session" + with pytest.raises(TransportReached): + await handler.chat_completion(model, "sys", "usr") + actual = seen[-1] + assert actual.url == native.url + assert actual.headers["Authorization"].startswith("AWS4-HMAC-SHA256 ") + assert f"/{region}/bedrock/aws4_request" in actual.headers["Authorization"] + assert "Credential=owned-key/" in actual.headers["Authorization"] + assert actual.headers["X-Amz-Security-Token"] == "owned-session" + assert handler._aws_active_creds == original_credentials + assert len(seen) == 4 + + +@pytest.mark.parametrize("auth", ("sigv4", "late_foreign", "captured_bearer")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.usefixtures("native_endpoint_runtime") +@pytest.mark.asyncio +async def test_mantle_http_auth(monkeypatch, auth, entrypoint): + settings = _make_settings({ + "aws.AWS_ACCESS_KEY_ID": "owned-key", "aws.AWS_SECRET_ACCESS_KEY": "owned-secret", + "aws.AWS_SESSION_TOKEN": "owned-session", "aws.AWS_REGION_NAME": "us-east-1", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + if auth == "captured_bearer": + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "owned-bearer") + handler = LiteLLMAIHandler() + if auth != "sigv4": + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "foreign-bearer") + monkeypatch.setattr(litellm, "api_key", "foreign-global") + seen = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + seen.append(request) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + model = "bedrock_mantle/openai.gpt-oss-120b" + with pytest.raises(TransportReached): + if entrypoint == "chat": + await handler.chat_completion(model, "sys", "usr") + else: + await handler.probe_completion(model) + assert len(seen) == 1 + request = seen[0] + assert str(request.url) == "https://bedrock-mantle.us-east-1.api.aws/v1/chat/completions" + authorization = request.headers["Authorization"] + if auth == "captured_bearer": + assert authorization == "Bearer owned-bearer" + else: + assert authorization.startswith("AWS4-HMAC-SHA256 ") + assert "Credential=owned-key/" in authorization + assert "/us-east-1/bedrock/aws4_request" in authorization + assert request.headers["X-Amz-Security-Token"] == "owned-session" + assert "foreign" not in str(request.headers) + assert litellm_handler.DUMMY_LITELLM_API_KEY not in authorization + import json + + body = json.loads(request.content) + assert not set(litellm_handler.BEDROCK_MANTLE_REQUEST_BODY_EXCLUDED_KEYS).intersection(body) + assert all(secret.encode() not in request.content for secret in ("owned-secret", "owned-session", "foreign")) + assert litellm_handler._bedrock_mantle_request_credentials.get() is None + assert litellm_handler._bedrock_mantle_block_bearer.get() is False + + +@pytest.mark.usefixtures("native_endpoint_runtime") +@pytest.mark.parametrize("auth", ("sigv4", "bearer", "bearer_only")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.parametrize(("selection", "environment", "expected_region", "expected_base"), ( + ("absent", {}, "us-east-1", None), + ("default", {"AWS_DEFAULT_REGION": "eu-west-1"}, "us-east-1", None), + ("default_late", {}, "us-east-1", None), + ("region_late", {}, "us-east-1", None), + ("region_name_late", {}, "us-east-1", None), + ("mantle_late", {}, "us-east-1", None), + ("region", {"AWS_DEFAULT_REGION": "eu-west-1", "AWS_REGION": "us-west-2"}, "us-west-2", None), + ("region_name", {"AWS_REGION": "us-west-2", "AWS_REGION_NAME": "eu-west-1"}, "eu-west-1", None), + ("mantle", {"AWS_REGION_NAME": "eu-west-1", "BEDROCK_MANTLE_REGION": "us-east-1"}, "us-east-1", None), + ("settings", {"AWS_DEFAULT_REGION": "eu-west-1"}, "us-west-2", None), + ("host", {"BEDROCK_MANTLE_API_BASE": "https://bedrock-mantle.eu-west-1.api.aws/v1"}, + "eu-west-1", "https://bedrock-mantle.eu-west-1.api.aws/v1"), + ("custom", {"BEDROCK_MANTLE_API_BASE": "https://mantle.example/v1", "AWS_REGION": "us-west-2"}, + "us-west-2", "https://mantle.example/v1"), +)) +@pytest.mark.asyncio +async def test_native_mantle_region(monkeypatch, auth, entrypoint, selection, environment, expected_region, expected_base): + for name in ("AWS_REGION_NAME", "AWS_REGION", "AWS_DEFAULT_REGION", "DEFAULT_REGION", "BEDROCK_MANTLE_REGION"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AWS_USE_IMDS", "false") + for name, value in environment.items(): + monkeypatch.setenv(name, value) + credentials = {"aws_access_key_id": "owned-key", "aws_secret_access_key": "owned-secret", + "aws_session_token": "owned-session"} + if auth != "bearer_only": + for name, value in credentials.items(): + monkeypatch.setenv(name.upper(), value) + native_params = dict(credentials) + if auth != "sigv4": + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "owned-bearer") + native_params = {"api_key": "owned-bearer"} + settings = _make_settings({"aws.AWS_REGION_NAME": expected_region} if selection == "settings" else {}) + if selection == "settings": + # The core deliberately supports settings-only region for Mantle. + native_params["aws_region_name"] = expected_region + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + seen = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + seen.append(request) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + model = "bedrock_mantle/openai.gpt-oss-120b" + with pytest.raises(TransportReached): + await litellm.acompletion(model=model, messages=[{"role": "user", "content": "usr"}], **native_params) + base = expected_base or f"https://bedrock-mantle.{expected_region}.api.aws/v1" + assert str(seen[0].url) == base + "/chat/completions" + + def assert_auth(request): + if auth != "sigv4": + assert request.headers["Authorization"] == "Bearer owned-bearer" + else: + assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256 ") + assert "Credential=owned-key/" in request.headers["Authorization"] + assert f"/{expected_region}/bedrock/aws4_request" in request.headers["Authorization"] + assert request.headers["X-Amz-Security-Token"] == "owned-session" + + assert_auth(seen[0]) + handler = LiteLLMAIHandler() + captured = dict(handler._aws_active_creds) + late_variable = { + "default_late": "AWS_DEFAULT_REGION", "region_late": "AWS_REGION", + "region_name_late": "AWS_REGION_NAME", "mantle_late": "BEDROCK_MANTLE_REGION", + }.get(selection) + if late_variable: + monkeypatch.setenv(late_variable, "eu-west-1") + with pytest.raises(TransportReached): + if entrypoint == "chat": + await handler.chat_completion(model, "sys", "usr") + else: + await handler.probe_completion(model) + assert len(seen) == 2 + assert seen[1].url == seen[0].url + assert_auth(seen[1]) + assert handler._aws_active_creds == captured + + +@pytest.mark.usefixtures("native_endpoint_runtime") +@pytest.mark.parametrize(("model", "selection"), ( + ("volcengine/doubao-pro-32k", "volcengine"), + ("volcengine/doubao-pro-32k", "both"), + ("volcengine/responses/doubao-pro-32k", "volcengine"), + ("volcengine/responses/doubao-pro-32k", "both"), + ("volcengine/responses/doubao-pro-32k", "ark"), + ("volcengine/responses/doubao-pro-32k", "ark_gateway"), + ("volcengine/responses/doubao-pro-32k", "gateway_late"), +)) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.asyncio +async def test_native_volcengine_key_precedence(monkeypatch, model, selection, entrypoint): + if selection in ("volcengine", "both"): + monkeypatch.setenv("VOLCENGINE_API_KEY", "owned-volcengine-key") + if selection in ("both", "ark", "ark_gateway"): + monkeypatch.setenv("ARK_API_KEY", "owned-ark-key") + settings = _make_settings({ + "OPENAI.KEY": "gateway-key", "OPENAI.API_BASE": "https://gateway.example/v1", + } if selection in ("ark_gateway", "gateway_late") else {}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + seen = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + seen.append(request) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + native_params = ( + {"api_key": "gateway-key", "api_base": "https://gateway.example/v1"} + if selection == "gateway_late" else {} + ) + with pytest.raises(TransportReached): + await litellm.acompletion(model=model, messages=[{"role": "user", "content": "usr"}], **native_params) + expected = "owned-volcengine-key" if selection in ("volcengine", "both") else "owned-ark-key" + if selection == "gateway_late": + expected = "gateway-key" + assert seen[0].headers["Authorization"] == f"Bearer {expected}" + assert seen[0].url.host == ("gateway.example" if selection == "gateway_late" else "ark.cn-beijing.volces.com") + handler = LiteLLMAIHandler() + monkeypatch.setenv("VOLCENGINE_API_KEY", "foreign-volcengine-key") + monkeypatch.setenv("ARK_API_KEY", "foreign-ark-key") + with pytest.raises(TransportReached): + if entrypoint == "chat": + await handler.chat_completion(model, "sys", "usr") + else: + await handler.probe_completion(model) + assert len(seen) == 2 + assert seen[1].url == seen[0].url + assert seen[1].headers.get_list("Authorization") == seen[0].headers.get_list("Authorization") + + +@pytest.mark.usefixtures("native_endpoint_runtime") +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.asyncio +async def test_volcengine_responses_rejects_late_ark_key(monkeypatch, entrypoint): + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings()) + handler = LiteLLMAIHandler() + monkeypatch.setenv("ARK_API_KEY", "foreign-ark-key") + seen = [] + + async def send(client, request, **kwargs): + seen.append(request) + raise AssertionError("A late ARK key must not reach HTTP transport") + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + with pytest.raises(ValueError, match="Refusing Volcengine ARK API key added after handler initialization"): + if entrypoint == "chat": + await handler.chat_completion("volcengine/responses/doubao-pro-32k", "sys", "usr") + else: + await handler.probe_completion("volcengine/responses/doubao-pro-32k") + assert not seen + + +@pytest.mark.usefixtures("native_endpoint_runtime") +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.asyncio +async def test_volcengine_chat_ignores_ark_only_key(monkeypatch, entrypoint): + monkeypatch.setenv("ARK_API_KEY", "owned-ark-key") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings()) + seen = [] + + async def send(client, request, **kwargs): + seen.append(request) + raise AssertionError("Ordinary chat must not authenticate with ARK_API_KEY") + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + model = "volcengine/doubao-pro-32k" + with pytest.raises(litellm.InternalServerError, match="Missing credentials"): + await litellm.acompletion(model=model, messages=[{"role": "user", "content": "usr"}]) + handler = LiteLLMAIHandler() + assert "api_key" not in handler._get_provider_request_params(model) + with pytest.raises(litellm.InternalServerError, match="Missing credentials"): + if entrypoint == "chat": + await handler.chat_completion(model, "sys", "usr") + else: + await handler.probe_completion(model) + assert not seen + + +@pytest.mark.usefixtures("native_endpoint_runtime") +@pytest.mark.parametrize("route", ("converse", "invoke", "implicit_invoke")) +@pytest.mark.parametrize("model_source", ("plain", "model_id", "model_arn", "region_path")) +@pytest.mark.parametrize("region_source", ( + "AWS_REGION_NAME", "AWS_REGION", "AWS_DEFAULT_REGION", "settings", "static", "imds", +)) +@pytest.mark.parametrize("auth", ("sigv4", "bearer")) +@pytest.mark.asyncio +async def test_native_bedrock_model_region_precedence(monkeypatch, model_source, region_source, auth, route): + if route != "converse" and model_source == "region_path": + pytest.skip("Native region/model syntax is only supported by Converse") + native_model = ( + "mistral.mistral-7b-instruct-v0:2" if route == "implicit_invoke" + else "anthropic.claude-3-haiku-20240307-v1:0" + ) + arn = f"arn:aws:bedrock:eu-west-1:123456789012:inference-profile/{native_model}" + model = f"bedrock/{native_model}" + overrides = {"litellm.model_id": arn} if model_source == "model_id" else {} + if model_source == "model_arn": + model = f"bedrock/{arn}" + elif model_source == "region_path": + model = f"bedrock/eu-west-1/{native_model}" + if route == "invoke": + model = model.replace("bedrock/", "bedrock/invoke/", 1) + for name in ("AWS_REGION_NAME", "AWS_REGION", "AWS_DEFAULT_REGION", "DEFAULT_REGION"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AWS_USE_IMDS", "true" if region_source == "imds" else "false") + credentials = {"aws_access_key_id": "owned-key", "aws_secret_access_key": "owned-secret", + "aws_session_token": "owned-session"} + for name, value in credentials.items(): + monkeypatch.setenv(name.upper(), value) + if region_source in ("settings", "static"): + overrides["aws.AWS_REGION_NAME"] = "us-east-1" + if region_source == "static": + overrides.update({f"aws.{name.upper()}": value for name, value in credentials.items()}) + else: + monkeypatch.setenv("AWS_REGION_NAME" if region_source == "imds" else region_source, "us-east-1") + if region_source == "imds": + session = MagicMock() + session.region_name = "us-east-1" + session.get_credentials.return_value.get_frozen_credentials.return_value = SimpleNamespace( + access_key="owned-key", secret_key="owned-secret", token="owned-session", + ) + monkeypatch.setattr("boto3.Session", lambda **kwargs: session) + native_params = dict(credentials) + if auth == "bearer": + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "owned-bearer") + native_params = {"api_key": "owned-bearer"} + if model_source == "model_id": + native_params["model_id"] = arn + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) + seen = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + seen.append(request) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + with monkeypatch.context() as native_environment: + if region_source in ("settings", "static"): + # The original handler exported the settings region, not a kwarg. + native_environment.setenv("AWS_REGION_NAME", "us-east-1") + with pytest.raises(TransportReached): + await litellm.acompletion(model=model, messages=[{"role": "user", "content": "usr"}], **native_params) + expected_region = ( + "eu-west-1" if model_source == "model_arn" or (route == "converse" and model_source != "plain") + else "us-east-1" + ) + + def assert_request(request): + assert request.url.host == f"bedrock-runtime.{expected_region}.amazonaws.com" + if auth == "bearer": + assert request.headers["Authorization"] == "Bearer owned-bearer" + else: + assert "Credential=owned-key/" in request.headers["Authorization"] + assert f"/{expected_region}/bedrock/aws4_request" in request.headers["Authorization"] + assert request.headers["X-Amz-Security-Token"] == "owned-session" + + assert len(seen) == 1 + assert_request(seen[0]) + handler = LiteLLMAIHandler() + captured = dict(handler._aws_active_creds) + with pytest.raises(TransportReached): + await handler.chat_completion(model, "sys", "usr") + assert len(seen) == 2 + assert seen[1].url == seen[0].url + assert_request(seen[1]) + assert handler._aws_active_creds == captured + + +@pytest.mark.parametrize( + ("api_base", "expected"), + ( + ("https://gateway.ai.cloudflare.com/v1/account/gateway/azure-openai/resource", True), + ("https://eu.gateway.ai.cloudflare.com/v1/account/gateway/azure-openai/resource", True), + ("https://owned.example/?next=gateway.ai.cloudflare.com", False), + ("https://owned.example/gateway.ai.cloudflare.com/v1/account", False), + ("https://gateway.ai.cloudflare.com.owned.example/v1", False), + ("https://owned.openai.azure.com", False), + (None, False), + ), +) +def test_cloudflare_gateway_matches_the_host_not_the_string(api_base, expected): + assert litellm_handler._is_cloudflare_gateway(api_base) is expected diff --git a/tests/unittest/test_litellm_chat_completion_core.py b/tests/unittest/test_litellm_chat_completion_core.py index 000ff62a5d..f9d3ec6a38 100644 --- a/tests/unittest/test_litellm_chat_completion_core.py +++ b/tests/unittest/test_litellm_chat_completion_core.py @@ -9,6 +9,22 @@ import pr_agent.algo.ai_handlers.litellm_ai_handler as litellm_handler +@pytest.fixture(autouse=True) +def isolate_aws_environment(monkeypatch): + environment_variables = set(litellm_handler.AWS_CREDENTIAL_CHAIN_ENV_VARS) | { + "AWS_USE_IMDS", + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_DEFAULT_REGION", + "AWS_BEARER_TOKEN_BEDROCK", + } + for variable in environment_variables: + monkeypatch.delenv(variable, raising=False) + + class FakeBox: def __init__(self, values=None, **attrs): self._values = values or {} @@ -31,7 +47,12 @@ def __init__(self, config_values=None, settings_values=None): model="gpt-4o", ) self.litellm = FakeBox() - self._settings_values = settings_values or {} + self._settings_values = { + "aws.AWS_ACCESS_KEY_ID": "test-access-key", + "aws.AWS_SECRET_ACCESS_KEY": "test-secret-key", + "aws.AWS_REGION_NAME": "us-east-1", + **(settings_values or {}), + } def get(self, key, default=None): return self._settings_values.get(key, default) @@ -64,6 +85,22 @@ async def test_chat_completion_passes_seed_when_temperature_is_zero(monkeypatch) assert mock_call.call_args.kwargs["seed"] == 123 +@pytest.mark.asyncio +async def test_claude_empty_system_prompt_uses_public_request_normalization(monkeypatch): + monkeypatch.setattr(litellm_handler, "get_settings", FakeSettings) + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + handler = litellm_handler.LiteLLMAIHandler() + expected_prompts = handler.normalize_request_prompts("claude-sonnet-4-5", "", "usr") + + await handler.chat_completion(model="claude-sonnet-4-5", system="", user="usr") + + messages = mock_call.call_args.kwargs["messages"] + assert expected_prompts == ("No system prompt provided", "usr") + assert (messages[0]["content"], messages[1]["content"]) == expected_prompts + + @pytest.mark.asyncio async def test_chat_completion_probes_images_off_loop_with_timeout(monkeypatch): monkeypatch.setattr(litellm_handler, "get_settings", FakeSettings) @@ -96,6 +133,31 @@ def fake_head(url, **kwargs): } +@pytest.mark.asyncio +async def test_chat_completion_dead_image_uses_current_help_link(monkeypatch): + monkeypatch.setattr(litellm_handler, "get_settings", FakeSettings) + monkeypatch.setattr( + litellm_handler.requests, + "head", + lambda *args, **kwargs: SimpleNamespace(status_code=404), + ) + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + handler = litellm_handler.LiteLLMAIHandler() + + response, finish_reason = await handler.chat_completion( + model="gpt-4o", + system="sys", + user="usr", + img_path="https://example.test/missing.png", + ) + + assert "https://docs.pr-agent.ai/tools/ask/#ask-on-images" in response + assert "pr-agent-docs.codium.ai" not in response + assert finish_reason == "error" + mock_call.assert_not_awaited() + + @pytest.mark.asyncio @pytest.mark.parametrize( ("model", "expected_model_id"), @@ -123,6 +185,22 @@ async def test_chat_completion_scopes_model_id_to_classic_bedrock(monkeypatch, m assert mock_call.call_args.kwargs["model_id"] == expected_model_id +@pytest.mark.asyncio +async def test_health_probe_uses_snapshotted_classic_bedrock_model_id(monkeypatch): + active_settings = FakeSettings(settings_values={"litellm.model_id": "profile-a"}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + handler = litellm_handler.LiteLLMAIHandler() + active_settings = FakeSettings(settings_values={"litellm.model_id": "profile-b"}) + completion = AsyncMock(return_value=_mock_response()) + + await handler.probe_completion( + "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + _completion=completion, + ) + + assert completion.call_args.kwargs["model_id"] == "profile-a" + + @pytest.mark.asyncio async def test_chat_completion_accumulates_usage_into_run_details(monkeypatch): from pr_agent.algo.run_details import init_run_details @@ -386,6 +464,69 @@ async def test_chat_completion_combines_prompts_for_user_message_only_models(mon assert messages == [{"role": "user", "content": "sys\n\n\nusr"}] +@pytest.mark.asyncio +async def test_chat_completion_keeps_image_for_user_message_only_models(monkeypatch): + monkeypatch.setattr(litellm_handler, "get_settings", FakeSettings) + monkeypatch.setattr( + litellm_handler.requests, + "head", + lambda *args, **kwargs: SimpleNamespace(status_code=200), + ) + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + handler = litellm_handler.LiteLLMAIHandler() + handler.user_message_only_models = ["user-only-model"] + + await handler.chat_completion( + model="user-only-model", + system="sys", + user="usr", + img_path="https://example.test/image.png", + ) + + messages = mock_call.call_args.kwargs["messages"] + assert messages == [ + { + "role": "user", + "content": [ + {"type": "text", "text": "sys\n\n\nusr"}, + {"type": "image_url", "image_url": {"url": "https://example.test/image.png"}}, + ], + } + ] + assert any(block.get("type") == "image_url" for block in messages[0]["content"]) + + +@pytest.mark.asyncio +async def test_chat_completion_keeps_image_for_custom_reasoning_models(monkeypatch): + settings = FakeSettings() + settings.config.custom_reasoning_model = True + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + monkeypatch.setattr( + litellm_handler.requests, + "head", + lambda *args, **kwargs: SimpleNamespace(status_code=200), + ) + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + handler = litellm_handler.LiteLLMAIHandler() + + await handler.chat_completion( + model="gpt-4o", + system="sys", + user="usr", + img_path="https://example.test/image.png", + ) + + messages = mock_call.call_args.kwargs["messages"] + assert messages[0]["content"][1] == { + "type": "image_url", + "image_url": {"url": "https://example.test/image.png"}, + } + + # Wiring tests for the retry knobs: the helpers (_should_retry_same_model, # _configured_client_retries) are unit-tested in test_litellm_retry_config.py, but those # tests keep passing when the @retry predicate or the kwargs pass-through in @@ -462,6 +603,7 @@ async def test_chat_completion_timeout_not_retried_same_model_when_disabled(monk @pytest.mark.asyncio async def test_get_completion_uses_streaming_for_required_models(): handler = litellm_handler.LiteLLMAIHandler.__new__(litellm_handler.LiteLLMAIHandler) + handler._sdk_header_defaults = {"organization": None, "project": None, "custom_headers": {}} handler.streaming_required_models = ["streaming-model"] with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call, \ @@ -487,6 +629,24 @@ async def test_get_completion_uses_streaming_for_required_models(): assert response_obj.dict()["choices"][0]["message"]["content"] == "streamed text" +@pytest.mark.parametrize("model", ("azure/qwq-plus", "azure/openai/qwq-plus")) +@pytest.mark.asyncio +async def test_get_completion_preserves_streaming_requirement_after_azure_routing(model): + handler = litellm_handler.LiteLLMAIHandler.__new__(litellm_handler.LiteLLMAIHandler) + handler._sdk_header_defaults = {"organization": None, "project": None, "custom_headers": {}} + handler.streaming_required_models = ["openai/qwq-plus"] + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call, \ + patch("pr_agent.algo.ai_handlers.litellm_ai_handler._handle_streaming_response", + new_callable=AsyncMock) as mock_stream: + mock_call.return_value = "stream" + mock_stream.return_value = ("streamed text", "stop", MagicMock()) + + await handler._get_completion(model=model, messages=[]) + + assert mock_call.call_args.kwargs["stream"] is True + + def _empty_content_response(finish_reason="stop"): mock = MagicMock() response = {"choices": [{"message": {"content": ""}, "finish_reason": finish_reason}]} @@ -513,6 +673,7 @@ async def test_get_completion_raises_on_empty_content_for_non_streaming_model(): @pytest.mark.asyncio async def test_get_completion_returns_non_empty_content_for_non_streaming_model(): handler = litellm_handler.LiteLLMAIHandler.__new__(litellm_handler.LiteLLMAIHandler) + handler._sdk_header_defaults = {"organization": None, "project": None, "custom_headers": {}} handler.streaming_required_models = [] with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: diff --git a/tests/unittest/test_litellm_claude_adaptive_thinking.py b/tests/unittest/test_litellm_claude_adaptive_thinking.py index ea2b92d4e0..fdcf13a230 100644 --- a/tests/unittest/test_litellm_claude_adaptive_thinking.py +++ b/tests/unittest/test_litellm_claude_adaptive_thinking.py @@ -43,10 +43,18 @@ def _restore_litellm_globals(): os.environ[name] = value -def _settings(reasoning_effort="medium", enabled=False, extended_enabled=False): +def _settings( + reasoning_effort="medium", + enabled=False, + extended_enabled=False, + extended_budget_tokens=2048, + extended_max_output_tokens=4096, +): flags = { "enable_claude_adaptive_thinking": enabled, "enable_claude_extended_thinking": extended_enabled, + "extended_thinking_budget_tokens": extended_budget_tokens, + "extended_thinking_max_output_tokens": extended_max_output_tokens, } config = SimpleNamespace( reasoning_effort=reasoning_effort, @@ -56,10 +64,17 @@ def _settings(reasoning_effort="medium", enabled=False, extended_enabled=False): verbosity_level=0, get=lambda key, default=None: flags.get(key, default), ) + # Bedrock requests now resolve credentials per request rather than from the + # process environment, so a bedrock model needs them supplied here. + aws = { + "aws.AWS_ACCESS_KEY_ID": "AKIAIOSFODNN7EXAMPLE", + "aws.AWS_SECRET_ACCESS_KEY": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "aws.AWS_REGION_NAME": "us-east-1", + } return SimpleNamespace( config=config, litellm=SimpleNamespace(get=lambda key, default=None: default), - get=lambda key, default=None: default, + get=lambda key, default=None: aws.get(key, default), ) @@ -73,6 +88,9 @@ def _response(): async def _run_completion(monkeypatch, model, reasoning_effort="medium", enabled=False, extended_enabled=False, extended_override=None): + # An ambient selector would be refused before the request credentials are read. + for variable in ("AWS_PROFILE_NAME", "AWS_ROLE_NAME"): + monkeypatch.delenv(variable, raising=False) monkeypatch.setattr( litellm_handler, "get_settings", @@ -126,6 +144,69 @@ async def test_enabled_adaptive_thinking_sends_anthropic_payload(monkeypatch): assert "reasoning_effort" not in kwargs +@pytest.mark.asyncio +async def test_adaptive_thinking_effort_is_isolated_from_later_settings(monkeypatch): + active_settings = _settings(reasoning_effort="low", enabled=True) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + handler = LiteLLMAIHandler() + active_settings = _settings(reasoning_effort="high", enabled=True) + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as completion: + completion.return_value = _response() + await handler.chat_completion(model="anthropic/claude-opus-4-8", system="sys", user="usr") + + assert completion.call_args.kwargs["output_config"] == {"effort": "low"} + + +@pytest.mark.asyncio +async def test_adaptive_thinking_enablement_is_isolated_from_later_settings(monkeypatch): + active_settings = _settings(reasoning_effort="low", enabled=False) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + handler = LiteLLMAIHandler() + active_settings = _settings(reasoning_effort="high", enabled=True) + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as completion: + completion.return_value = _response() + await handler.chat_completion(model="anthropic/claude-opus-4-8", system="sys", user="usr") + + assert "thinking" not in completion.call_args.kwargs + assert "output_config" not in completion.call_args.kwargs + + +@pytest.mark.asyncio +async def test_extended_thinking_limits_are_isolated_from_later_settings(monkeypatch): + model = "anthropic/claude-opus-4-6" + active_settings = _settings( + extended_enabled=True, + extended_budget_tokens=1024, + extended_max_output_tokens=2048, + ) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + handler = LiteLLMAIHandler() + handler.claude_extended_thinking_models = [model] + active_settings = _settings( + extended_enabled=True, + extended_budget_tokens=4096, + extended_max_output_tokens=8192, + ) + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as completion: + completion.return_value = _response() + await handler.chat_completion(model=model, system="sys", user="usr") + + assert completion.call_args.kwargs["thinking"] == {"type": "enabled", "budget_tokens": 1024} + assert completion.call_args.kwargs["max_tokens"] == 2048 + + @pytest.mark.asyncio async def test_adaptive_thinking_accepts_max_effort_without_enum_dependency(monkeypatch): kwargs = await _run_completion( @@ -207,3 +288,78 @@ async def test_non_adaptive_model_in_extended_override_still_gets_extended_think assert kwargs["thinking"]["type"] == "enabled" assert "budget_tokens" in kwargs["thinking"] + + +@pytest.mark.asyncio +async def test_opaque_arn_with_adaptive_enabled_warns_and_skips_payload(monkeypatch): + """An opaque application inference profile ARN carries no model name, so the adaptive + payload must be skipped and a warning logged instead of failing silently (see #3216).""" + logger = MagicMock() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.get_logger", + return_value=logger): + kwargs = await _run_completion( + monkeypatch, + "bedrock/converse/arn:aws:bedrock:eu-central-1:000000000000:application-inference-profile/abc123def456", + enabled=True, + ) + + assert "thinking" not in kwargs + assert "output_config" not in kwargs + logger.warning.assert_called_once() + assert "abc123def456" in logger.warning.call_args.args[0] + assert "litellm.model_id" in logger.warning.call_args.args[0] + + +@pytest.mark.asyncio +async def test_family_embedding_arn_still_gets_adaptive_payload(monkeypatch): + """An ARN that embeds the model family normalises to a matching id, so adaptive thinking + applies; the miss is specific to opaque suffixes (see #3216).""" + kwargs = await _run_completion( + monkeypatch, + "bedrock/converse/arn:aws:bedrock:eu-central-1:000000000000:inference-profile/us.anthropic.claude-sonnet-5", + enabled=True, + ) + + assert kwargs["thinking"] == {"type": "adaptive"} + assert kwargs["output_config"] == {"effort": "medium"} + + +@pytest.mark.asyncio +async def test_extended_enabled_with_no_matching_model_warns(monkeypatch): + """Extended thinking enabled with an empty override list leaves no matching model, so the + warning fires even though the adaptive model gate is not involved (see #3216).""" + logger = MagicMock() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.get_logger", + return_value=logger): + kwargs = await _run_completion( + monkeypatch, + "anthropic/claude-opus-4-6", + enabled=False, + extended_enabled=True, + extended_override=[], + ) + + assert "thinking" not in kwargs + logger.warning.assert_called_once() + assert "opus-4-6" in logger.warning.call_args.args[0] + + +@pytest.mark.asyncio +async def test_non_arn_model_warns_without_bedrock_advice(monkeypatch): + """The Bedrock remedy must stay gated to Bedrock ids: a plain provider id that reaches the + warning branch only gets the generic message, never the litellm.model_id ARN advice.""" + logger = MagicMock() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.get_logger", + return_value=logger): + kwargs = await _run_completion( + monkeypatch, + "openai/gpt-4o", + enabled=True, + ) + + assert "thinking" not in kwargs + logger.warning.assert_called_once() + message = logger.warning.call_args.args[0] + assert "gpt-4o" in message + assert "litellm.model_id" not in message + assert "arn" not in message diff --git a/tests/unittest/test_litellm_custom_provider.py b/tests/unittest/test_litellm_custom_provider.py index e8be216224..3313c79b35 100644 --- a/tests/unittest/test_litellm_custom_provider.py +++ b/tests/unittest/test_litellm_custom_provider.py @@ -1,5 +1,8 @@ +import asyncio +import inspect from unittest.mock import AsyncMock, patch +import openai import pytest import pr_agent.algo.ai_handlers.litellm_ai_handler as litellm_handler @@ -10,6 +13,9 @@ def create_mock_settings( custom_llm_provider=None, force_streaming_custom_llm_provider="openai", force_streaming_api_base_substrings=None, + openai_key=None, + anthropic_key=None, + deployment_id=None, ): if force_streaming_api_base_substrings is None: force_streaming_api_base_substrings = ["snowflakecomputing.com"] @@ -28,6 +34,9 @@ def get_value(key, default=None): "litellm.force_streaming_custom_llm_provider": force_streaming_custom_llm_provider, "LITELLM.FORCE_STREAMING_API_BASE_SUBSTRINGS": force_streaming_api_base_substrings, "litellm.force_streaming_api_base_substrings": force_streaming_api_base_substrings, + "OPENAI.KEY": openai_key, + "ANTHROPIC.KEY": anthropic_key, + "OPENAI.DEPLOYMENT_ID": deployment_id, } return values.get(key, default) @@ -41,6 +50,7 @@ def get_value(key, default=None): { "ai_timeout": 120, "custom_reasoning_model": False, + "reasoning_effort": None, "verbosity_level": 0, "get": lambda self, key, default=None: default, }, @@ -63,11 +73,102 @@ def dict(self): return MockCompletionResponse(response_payload) +@pytest.fixture(autouse=True) +def restore_anthropic_bridge(): + anthropic_get_api_key = inspect.getattr_static(litellm_handler.AnthropicModelInfo, "get_api_key") + anthropic_get_auth_token = inspect.getattr_static(litellm_handler.AnthropicModelInfo, "get_auth_token") + module_get_api_key = litellm_handler._anthropic_get_api_key + module_get_auth_token = litellm_handler._anthropic_get_auth_token + yield + litellm_handler.AnthropicModelInfo.get_api_key = anthropic_get_api_key + litellm_handler.AnthropicModelInfo.get_auth_token = anthropic_get_auth_token + litellm_handler._anthropic_get_api_key = module_get_api_key + litellm_handler._anthropic_get_auth_token = module_get_auth_token + + +@pytest.mark.parametrize("model", ("claude-sonnet-4-5", "gpt-5.1")) @pytest.mark.asyncio -async def test_custom_llm_provider_is_forwarded_without_rewriting_model(monkeypatch): +async def test_custom_llm_provider_is_forwarded_without_rewriting_model(monkeypatch, model): fake_settings = create_mock_settings(" OpenAI ") monkeypatch.setattr(litellm_handler, "get_settings", lambda: fake_settings) + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_completion: + mock_completion.return_value = create_mock_acompletion_response() + + handler = LiteLLMAIHandler() + handler.azure = True + await handler.chat_completion( + model=model, + system="test system", + user="test user", + ) + + call_kwargs = mock_completion.call_args[1] + assert call_kwargs["model"] == model + assert call_kwargs["custom_llm_provider"] == "openai" + + +@pytest.mark.parametrize("custom_llm_provider", ("azure", "azure_text")) +@pytest.mark.asyncio +async def test_custom_azure_text_model_preserves_deployment_routing(monkeypatch, custom_llm_provider): + fake_settings = create_mock_settings(custom_llm_provider, deployment_id="custom-deployment") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: fake_settings) + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_completion: + mock_completion.return_value = create_mock_acompletion_response() + handler = LiteLLMAIHandler() + handler.azure = True + + await handler.chat_completion( + model="azure_text/gpt-35-turbo-instruct", + system="test system", + user="test user", + ) + + call_kwargs = mock_completion.call_args.kwargs + assert call_kwargs["model"] == "azure_text/custom-deployment" + assert call_kwargs["custom_llm_provider"] == custom_llm_provider + assert "deployment_id" not in call_kwargs + + +@pytest.mark.asyncio +async def test_bare_custom_azure_text_model_uses_deployment_as_model(monkeypatch): + fake_settings = create_mock_settings("azure_text", deployment_id="custom-deployment") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: fake_settings) + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_completion: + mock_completion.return_value = create_mock_acompletion_response() + await LiteLLMAIHandler().chat_completion( + model="gpt-35-turbo-instruct", + system="test system", + user="test user", + ) + + call_kwargs = mock_completion.call_args.kwargs + assert call_kwargs["model"] == "azure_text/custom-deployment" + assert call_kwargs["custom_llm_provider"] == "azure_text" + assert "deployment_id" not in call_kwargs + + +@pytest.mark.asyncio +async def test_custom_llm_provider_selects_request_credentials(monkeypatch): + fake_settings = create_mock_settings( + "openai", + openai_key="request-openai-key", + anthropic_key="unrelated-anthropic-key", + ) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: fake_settings) + monkeypatch.setenv("OPENAI_API_KEY", "ambient-openai-key") + with patch( "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock, @@ -82,10 +183,113 @@ async def test_custom_llm_provider_is_forwarded_without_rewriting_model(monkeypa ) call_kwargs = mock_completion.call_args[1] - assert call_kwargs["model"] == "claude-sonnet-4-5" + assert call_kwargs["api_key"] == "request-openai-key" assert call_kwargs["custom_llm_provider"] == "openai" +@pytest.mark.asyncio +async def test_custom_llm_provider_is_isolated_between_handlers(monkeypatch): + openai_settings = create_mock_settings("openai", openai_key="openai-key") + anthropic_settings = create_mock_settings("anthropic", anthropic_key="anthropic-key") + active_settings = openai_settings + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + + openai_handler = LiteLLMAIHandler() + active_settings = anthropic_settings + anthropic_handler = LiteLLMAIHandler() + + calls = {} + both_started = asyncio.Event() + started = 0 + + async def completion(**kwargs): + nonlocal started + calls[kwargs["messages"][1]["content"]] = kwargs + started += 1 + if started == 2: + both_started.set() + await asyncio.wait_for(both_started.wait(), timeout=5) + return create_mock_acompletion_response() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=completion): + await asyncio.gather( + openai_handler.chat_completion(model="hosted-model", system="sys", user="openai"), + anthropic_handler.chat_completion(model="hosted-model", system="sys", user="anthropic"), + ) + + assert calls["openai"]["custom_llm_provider"] == "openai" + assert calls["openai"]["api_key"] == "openai-key" + assert calls["anthropic"]["custom_llm_provider"] == "anthropic" + assert calls["anthropic"]["api_key"] == "anthropic-key" + + +@pytest.mark.asyncio +async def test_probe_uses_custom_llm_provider_captured_at_handler_creation(monkeypatch): + openai_settings = create_mock_settings("openai", openai_key="openai-key") + anthropic_settings = create_mock_settings("anthropic", anthropic_key="anthropic-key") + active_settings = openai_settings + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + handler = LiteLLMAIHandler() + active_settings = anthropic_settings + completion = AsyncMock(return_value=create_mock_acompletion_response()) + + await handler.probe_completion("hosted-model", _completion=completion) + + assert completion.call_args.kwargs["custom_llm_provider"] == "openai" + assert completion.call_args.kwargs["api_key"] == "openai-key" + + +@pytest.mark.parametrize("experimental_chat_handler", (False, True)) +@pytest.mark.asyncio +async def test_text_completion_openai_custom_provider_blocks_global_headers(monkeypatch, experimental_chat_handler): + fake_settings = create_mock_settings("text-completion-openai", openai_key="request-openai-key") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: fake_settings) + monkeypatch.setattr(litellm_handler.litellm, "headers", {"Authorization": "Bearer stale-key"}) + monkeypatch.delenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", raising=False) + if experimental_chat_handler: + monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_completion: + mock_completion.return_value = create_mock_acompletion_response() + + await LiteLLMAIHandler().chat_completion( + model="gpt-3.5-turbo-instruct", + system="test system", + user="test user", + ) + + call_kwargs = mock_completion.call_args.kwargs + assert call_kwargs["api_key"] == "request-openai-key" + assert call_kwargs["custom_llm_provider"] == "text-completion-openai" + assert set(call_kwargs["headers"]) == {"OpenAI-Organization", "OpenAI-Project"} + assert isinstance(call_kwargs["headers"]["OpenAI-Organization"], openai.Omit) + assert isinstance(call_kwargs["headers"]["OpenAI-Project"], openai.Omit) + + +@pytest.mark.asyncio +async def test_custom_anthropic_provider_uses_request_local_auth_token(monkeypatch): + fake_settings = create_mock_settings("anthropic") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: fake_settings) + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "request-auth-token") + + async def completion(**kwargs): + assert kwargs["api_key"] == litellm_handler.DUMMY_LITELLM_API_KEY + assert litellm_handler.AnthropicModelInfo.get_api_key(kwargs["api_key"]) is None + assert litellm_handler.AnthropicModelInfo.get_auth_token() == "request-auth-token" + return create_mock_acompletion_response() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=completion): + await LiteLLMAIHandler().chat_completion( + model="hosted-model", + system="test system", + user="test user", + ) + + @pytest.mark.asyncio async def test_custom_llm_provider_is_omitted_when_unset(monkeypatch): fake_settings = create_mock_settings() diff --git a/tests/unittest/test_litellm_databricks_provider.py b/tests/unittest/test_litellm_databricks_provider.py index 7c15ee7687..3ef0b1e6c1 100644 --- a/tests/unittest/test_litellm_databricks_provider.py +++ b/tests/unittest/test_litellm_databricks_provider.py @@ -1,33 +1,22 @@ -""" -Tests for Databricks provider wiring in LiteLLMAIHandler.__init__. +"""Tests for request-local Databricks provider wiring.""" -Verifies that DATABRICKS.API_KEY / DATABRICKS.API_BASE settings are exported to -the env vars LiteLLM's Databricks provider reads (DATABRICKS_API_KEY / -DATABRICKS_API_BASE), and that nothing is exported when they are unset. -""" +import asyncio +import atexit import os +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch +import httpx import litellm import pytest +import requests +from litellm.llms.databricks.common_utils import DatabricksBase import pr_agent.algo.ai_handlers.litellm_ai_handler as litellm_handler -# Env vars LiteLLMAIHandler.__init__ branches on — clear them so the handler -# under test isn't influenced by (or leaking into) the runner environment. -_ISOLATED_ENV = ( - "DATABRICKS_API_KEY", - "DATABRICKS_API_BASE", - "OPENAI_API_KEY", - "AWS_USE_IMDS", - "AWS_ACCESS_KEY_ID", - "AWS_SECRET_ACCESS_KEY", - "AWS_SESSION_TOKEN", - "AWS_REGION_NAME", -) - def _make_settings(overrides): - """Minimal settings whose top-level .get() returns the provided overrides.""" return type("Settings", (), { "config": type("Config", (), { "reasoning_effort": None, @@ -45,50 +34,346 @@ def _make_settings(overrides): })() +def _mock_response(): + mock = MagicMock() + mock.__getitem__ = lambda self, key: { + "choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}] + }[key] + mock.dict.return_value = {"choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}]} + return mock + + @pytest.fixture(autouse=True) -def _isolate_env(monkeypatch): - for var in _ISOLATED_ENV: - monkeypatch.delenv(var, raising=False) - # LiteLLMAIHandler.__init__ mutates the global litellm.api_key (sets the dummy - # fallback when OPENAI_API_KEY is absent, as it is here). Snapshot it so this - # file can't leak that global into order-dependent later tests. - saved_api_key = litellm.api_key - yield - litellm.api_key = saved_api_key - # Drop anything the handler wrote so it can't leak into other tests; - # monkeypatch then restores any pre-existing originals. - for var in ("DATABRICKS_API_KEY", "DATABRICKS_API_BASE"): - os.environ.pop(var, None) - - -def test_databricks_env_vars_exported_from_settings(monkeypatch): +def isolate_env(monkeypatch): + for variable in ( + "DATABRICKS_API_KEY", "DATABRICKS_API_BASE", "DATABRICKS_CLIENT_ID", "DATABRICKS_CLIENT_SECRET", + "AWS_USE_IMDS", "OPENAI_API_KEY", + ): + monkeypatch.delenv(variable, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setattr(litellm, "databricks_key", None) + validator = DatabricksBase.databricks_validate_environment + monkeypatch.setattr( + DatabricksBase, "databricks_validate_environment", + getattr(validator, "_pr_agent_original_databricks_validate_environment", validator), + ) + + +@pytest.mark.asyncio +async def test_databricks_settings_are_forwarded_without_exporting_env(monkeypatch): overrides = { "DATABRICKS.API_KEY": "dapi-test-123", "DATABRICKS.API_BASE": "https://adb-1234.azuredatabricks.net/serving-endpoints", } monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) - litellm_handler.LiteLLMAIHandler() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + handler = litellm_handler.LiteLLMAIHandler() + await handler.chat_completion(model="databricks/endpoint", system="sys", user="usr") + + assert mock_call.call_args.kwargs["api_key"] == "dapi-test-123" + assert mock_call.call_args.kwargs["api_base"] == overrides["DATABRICKS.API_BASE"] + assert "DATABRICKS_API_KEY" not in os.environ + assert "DATABRICKS_API_BASE" not in os.environ + + +@pytest.mark.asyncio +async def test_databricks_api_base_is_optional(monkeypatch): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"DATABRICKS.API_KEY": "dapi-only-key"}), + ) + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + handler = litellm_handler.LiteLLMAIHandler() + await handler.chat_completion(model="databricks/endpoint", system="sys", user="usr") + + assert mock_call.call_args.kwargs["api_key"] == "dapi-only-key" + assert "api_base" not in mock_call.call_args.kwargs + + +@pytest.mark.asyncio +async def test_databricks_native_env_is_forwarded_without_openai_gateway(monkeypatch): + monkeypatch.setenv("DATABRICKS_API_KEY", "native-databricks-key") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://databricks.example") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings({"OPENAI.API_BASE": "https://gateway.example/v1"}), + ) - assert os.environ["DATABRICKS_API_KEY"] == "dapi-test-123" - assert os.environ["DATABRICKS_API_BASE"] == "https://adb-1234.azuredatabricks.net/serving-endpoints" + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + handler = litellm_handler.LiteLLMAIHandler() + await handler.chat_completion(model="databricks/endpoint", system="sys", user="usr") + assert mock_call.call_args.kwargs["api_key"] == "native-databricks-key" + assert mock_call.call_args.kwargs["api_base"] == "https://databricks.example" -def test_databricks_env_vars_absent_when_unset(monkeypatch): + +@pytest.mark.asyncio +async def test_databricks_native_endpoint_is_frozen_with_its_key(monkeypatch): monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({})) + monkeypatch.setenv("DATABRICKS_API_KEY", "native-databricks-key") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://tenant-a.example") + handler = litellm_handler.LiteLLMAIHandler() + monkeypatch.setenv("DATABRICKS_API_KEY", "another-request-key") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://tenant-b.example") - litellm_handler.LiteLLMAIHandler() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + await handler.chat_completion(model="databricks/endpoint", system="sys", user="usr") - assert "DATABRICKS_API_KEY" not in os.environ - assert "DATABRICKS_API_BASE" not in os.environ + assert mock_call.call_args.kwargs["api_key"] == "native-databricks-key" + assert mock_call.call_args.kwargs["api_base"] == "https://tenant-a.example" -def test_databricks_api_base_optional(monkeypatch): - """API base is optional (a workspace default may be configured elsewhere).""" - overrides = {"DATABRICKS.API_KEY": "dapi-only-key"} +@pytest.mark.asyncio +async def test_databricks_does_not_receive_another_provider_credentials(monkeypatch): + overrides = { + "OPENROUTER.KEY": "openrouter-key", + "OPENROUTER.API_BASE": "https://openrouter.example/v1", + } monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides)) - litellm_handler.LiteLLMAIHandler() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + handler = litellm_handler.LiteLLMAIHandler() + await handler.chat_completion(model="databricks/endpoint", system="sys", user="usr") - assert os.environ["DATABRICKS_API_KEY"] == "dapi-only-key" - assert "DATABRICKS_API_BASE" not in os.environ + assert "api_key" not in mock_call.call_args.kwargs + assert "api_base" not in mock_call.call_args.kwargs + + +@pytest.mark.parametrize("entrypoint", ("chat_completion", "probe_completion")) +@pytest.mark.parametrize("key_source", ("global_before", "global_after", "provider_global", "captured_pat", "literal_dummy")) +@pytest.mark.asyncio +async def test_databricks_native_oauth_is_not_overwritten_by_foreign_keys(monkeypatch, entrypoint, key_source): + from litellm.litellm_core_utils import logging_worker + + endpoint = "https://databricks.example/serving-endpoints" + monkeypatch.setenv("DATABRICKS_API_BASE", endpoint) + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "request-client") + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "request-secret") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({})) + expected_token = "request-oauth-token" + if key_source in ("captured_pat", "literal_dummy"): + expected_token = "request-pat" if key_source == "captured_pat" else litellm_handler.DUMMY_LITELLM_API_KEY + monkeypatch.setenv("DATABRICKS_API_KEY", expected_token) + if key_source == "global_before": + monkeypatch.setattr(litellm, "api_key", "foreign-global-key") + if key_source == "provider_global": + monkeypatch.setattr(litellm, "databricks_key", "foreign-provider-key") + handler = litellm_handler.LiteLLMAIHandler() + if key_source in ("global_after", "captured_pat", "literal_dummy"): + monkeypatch.setattr(litellm, "api_key", "foreign-global-key") + monkeypatch.setenv("DATABRICKS_API_KEY", "late-pat") + oauth_calls, inference = [], [] + + def oauth_post(url, **kwargs): + assert url == "https://databricks.example/oidc/v1/token" + assert kwargs["auth"] == ("request-client", "request-secret") + oauth_calls.append(url) + response = requests.Response() + response.status_code = 200 + response._content = b'{"access_token": "request-oauth-token"}' + return response + + async def send(client, request, **kwargs): + assert str(request.url) == endpoint + "/chat/completions" + inference.append(request.headers.get("authorization")) + return httpx.Response(200, request=request, json={ + "id": "test", "object": "chat.completion", "created": 0, "model": "endpoint", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + + monkeypatch.setattr(requests, "post", oauth_post) + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + if entrypoint == "chat_completion": + assert await handler.chat_completion("databricks/endpoint", "sys", "usr") == ("ok", "stop") + else: + await handler.probe_completion("databricks/endpoint") + finally: + await asyncio.sleep(0) + await worker.flush() + await worker.stop() + assert len(oauth_calls) == 1 + assert inference == [f"Bearer {expected_token}"] + + +def _native_validation(api_key, api_base="https://databricks.example/serving-endpoints", positional=False): + validator = DatabricksBase().databricks_validate_environment + if positional: + return validator(api_key, api_base, "chat_completions", False, None) + return validator( + api_key=api_key, api_base=api_base, endpoint_type="chat_completions", custom_endpoint=False, headers=None, + ) + + +@pytest.fixture +def native_oauth(monkeypatch): + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "request-client") + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "request-secret") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({})) + token = MagicMock(return_value="request-oauth-token") + monkeypatch.setattr(DatabricksBase, "_get_oauth_m2m_token", token) + return token + + +@pytest.mark.parametrize("positional", (False, True)) +@pytest.mark.asyncio +async def test_databricks_bridge_only_translates_scoped_placeholder(native_oauth, positional): + handler = litellm_handler.LiteLLMAIHandler() + dummy = litellm_handler.DUMMY_LITELLM_API_KEY + + async def completion(**kwargs): + assert litellm_handler._databricks_request_keyless.get() is True + assert _native_validation("actual-pat", positional=positional)[1]["Authorization"] == "Bearer actual-pat" + return _native_validation(kwargs["api_key"], positional=positional) + + result = await handler._acompletion(_completion=completion, model="databricks/endpoint", api_key=dummy) + assert result[1]["Authorization"] == "Bearer request-oauth-token" + assert litellm_handler._databricks_request_keyless.get() is False + assert _native_validation(dummy, positional=positional)[1]["Authorization"] == f"Bearer {dummy}" + installed = DatabricksBase.databricks_validate_environment + litellm_handler._install_databricks_keyless_bridge() + assert DatabricksBase.databricks_validate_environment is installed + + +@pytest.mark.asyncio +async def test_databricks_concurrent_and_nested_literal_pat_keeps_context(native_oauth, monkeypatch): + keyless = litellm_handler.LiteLLMAIHandler() + dummy = litellm_handler.DUMMY_LITELLM_API_KEY + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({"DATABRICKS.API_KEY": dummy})) + pat = litellm_handler.LiteLLMAIHandler() + entered, release = asyncio.Event(), asyncio.Event() + + async def pat_completion(**kwargs): + assert litellm_handler._databricks_request_keyless.get() is False + return _native_validation(kwargs["api_key"])[1]["Authorization"] + + async def keyless_completion(**kwargs): + entered.set() + await release.wait() + assert await pat._acompletion( + _completion=pat_completion, model="databricks/endpoint", api_key=dummy, + ) == f"Bearer {dummy}" + assert litellm_handler._databricks_request_keyless.get() is True + return _native_validation(kwargs["api_key"])[1]["Authorization"] + + task = asyncio.create_task(keyless._acompletion( + _completion=keyless_completion, model="databricks/endpoint", api_key=dummy, + )) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + assert litellm_handler._databricks_request_keyless.get() is False + assert _native_validation(dummy)[1]["Authorization"] == f"Bearer {dummy}" + assert await pat._acompletion( + _completion=pat_completion, model="databricks/endpoint", api_key=dummy, + ) == f"Bearer {dummy}" + finally: + release.set() + result = await asyncio.wait_for(task, timeout=5) + assert result == "Bearer request-oauth-token" + assert litellm_handler._databricks_request_keyless.get() is False + + +@pytest.mark.parametrize("cancel", (False, True), ids=("failure", "cancellation")) +@pytest.mark.asyncio +async def test_databricks_context_resets_after_interrupted_completion(native_oauth, cancel): + handler = litellm_handler.LiteLLMAIHandler() + dummy = litellm_handler.DUMMY_LITELLM_API_KEY + entered = asyncio.Event() + + async def completion(**kwargs): + assert _native_validation(kwargs["api_key"])[1]["Authorization"] == "Bearer request-oauth-token" + entered.set() + if cancel: + await asyncio.Future() + raise RuntimeError("completion failed") + + async def request(): + try: + await handler._acompletion(_completion=completion, model="databricks/endpoint", api_key=dummy) + finally: + # Check the task that owned the context, not only its unaffected parent. + assert litellm_handler._databricks_request_keyless.get() is False + assert _native_validation(dummy)[1]["Authorization"] == f"Bearer {dummy}" + + task = asyncio.create_task(request()) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + finally: + if cancel or not entered.is_set(): + task.cancel() + with pytest.raises(asyncio.CancelledError if cancel else RuntimeError): + await asyncio.wait_for(task, timeout=5) + assert litellm_handler._databricks_request_keyless.get() is False + + +@pytest.mark.parametrize("interface", (None, lambda self: None, lambda self, api_key, /: None), + ids=("missing", "missing-api-key", "positional-only")) +@pytest.mark.asyncio +async def test_databricks_incompatible_interface_fails_before_completion(monkeypatch, interface): + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({})) + handler = litellm_handler.LiteLLMAIHandler() + monkeypatch.setattr(DatabricksBase, "databricks_validate_environment", interface) + completion = AsyncMock() + with pytest.raises(RuntimeError, match="Databricks authentication interface is incompatible"): + await handler._acompletion( + _completion=completion, model="databricks/endpoint", api_key=litellm_handler.DUMMY_LITELLM_API_KEY, + ) + completion.assert_not_awaited() + assert litellm_handler._databricks_request_keyless.get() is False + await handler._acompletion(_completion=completion, model="databricks/endpoint", api_key="actual-pat") + completion.assert_awaited_once() + + +@pytest.mark.parametrize("api_base", (None, "https://captured-endpoint.example/serving-endpoints")) +@pytest.mark.asyncio +async def test_databricks_optional_sdk_boundary_preserves_native_endpoint_contract(monkeypatch, api_base): + # Exercise LiteLLM's optional SDK boundary, not any installed SDK implementation. + sdk = ModuleType("databricks.sdk") + authenticate = MagicMock(return_value={"Authorization": "Bearer sdk-token"}) + sdk.WorkspaceClient = MagicMock(return_value=SimpleNamespace(config=SimpleNamespace( + host="https://sdk-host.example", authenticate=authenticate, + ))) + sdk.useragent = SimpleNamespace(with_partner=MagicMock()) + package = ModuleType("databricks") + package.__path__ = [] + package.sdk = sdk + monkeypatch.setitem(sys.modules, "databricks", package) + monkeypatch.setitem(sys.modules, "databricks.sdk", sdk) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({})) + handler = litellm_handler.LiteLLMAIHandler() + monkeypatch.setattr(litellm, "api_key", "foreign-global-key") + monkeypatch.setattr(litellm, "databricks_key", "foreign-provider-key") + monkeypatch.setenv("DATABRICKS_API_KEY", "late-pat") + dummy = litellm_handler.DUMMY_LITELLM_API_KEY + + async def completion(**kwargs): + return _native_validation(kwargs["api_key"], api_base=api_base) + + endpoint, headers = await handler._acompletion( + _completion=completion, model="databricks/endpoint", api_key=dummy, + ) + assert endpoint == (api_base or "https://sdk-host.example/serving-endpoints") + "/chat/completions" + assert headers["Authorization"] == "Bearer sdk-token" + sdk.WorkspaceClient.assert_called_once_with() + sdk.useragent.with_partner.assert_called_once_with("litellm") + authenticate.assert_called_once_with() + assert litellm.api_key == "foreign-global-key" + assert litellm.databricks_key == "foreign-provider-key" + assert os.environ["DATABRICKS_API_KEY"] == "late-pat" + assert litellm_handler._databricks_request_keyless.get() is False + assert _native_validation(dummy)[1]["Authorization"] == f"Bearer {dummy}" + authenticate.assert_called_once_with() diff --git a/tests/unittest/test_litellm_imds.py b/tests/unittest/test_litellm_imds.py index d916911540..a223b46ab3 100644 --- a/tests/unittest/test_litellm_imds.py +++ b/tests/unittest/test_litellm_imds.py @@ -1,41 +1,24 @@ -""" -Tests for AWS_USE_IMDS ambient credential support in LiteLLMAIHandler. - -Covers: - - Credentials resolved via boto3 and written to os.environ - - AWS_SESSION_TOKEN set/cleared correctly - - Region auto-resolved from boto3 when not configured - - Static keys stashed for fallback (including session token) - - boto3 failure falls through gracefully (no crash) - - No boto3 call when AWS_USE_IMDS is absent - - _refresh_aws_imds_credentials called before each Bedrock chat_completion - - Fallback to static keys on Bedrock API failure - - _activate_static_aws_fallback correctly restores/clears session token -""" +"""Tests for request-local AWS credentials used by LiteLLM Bedrock calls.""" + +import asyncio +import atexit import os +from datetime import UTC, datetime, timedelta from unittest.mock import AsyncMock, MagicMock, patch +from urllib.parse import unquote +import httpx +import litellm import openai import pytest -from botocore.exceptions import ClientError, CredentialRetrievalError -from tenacity import RetryError +from botocore.exceptions import ClientError, CredentialRetrievalError, ProfileNotFound import pr_agent.algo.ai_handlers.litellm_ai_handler as litellm_handler from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - -def _base_settings(extra_get=None): - """Minimal settings that satisfy __init__ and chat_completion.""" - def _get(self, key, default=None): - if extra_get: - v = extra_get(key) - if v is not None: - return v - return default +def _base_settings(overrides=None): + overrides = overrides or {} return type("Settings", (), { "config": type("Config", (), { "reasoning_effort": None, @@ -49,46 +32,23 @@ def _get(self, key, default=None): "litellm": type("LiteLLM", (), { "get": lambda self, key, default=None: default, })(), - "get": _get, + "get": lambda self, key, default=None: overrides.get(key, default), })() -def _mock_acompletion_response(usage=None): - mock = MagicMock() - response = {"choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}]} - if usage is not None: - response["usage"] = usage - mock.__getitem__ = lambda self, key: response[key] - mock.dict.return_value = response - return mock - - -def _static_aws_settings(session_token=None): - """Settings object with static AWS credentials configured.""" - keys = { - "aws.AWS_ACCESS_KEY_ID": "STATICKEY", - "aws.AWS_SECRET_ACCESS_KEY": "STATICSECRET", +def _static_aws_settings(prefix="STATIC", session_token=None, overrides=None): + overrides = { + "aws.AWS_ACCESS_KEY_ID": f"{prefix}-KEY", + "aws.AWS_SECRET_ACCESS_KEY": f"{prefix}-SECRET", "aws.AWS_REGION_NAME": "us-east-1", + **(overrides or {}), } if session_token: - keys["aws.AWS_SESSION_TOKEN"] = session_token - settings = _base_settings(extra_get=lambda key: keys.get(key)) - aws_attrs = { - "AWS_ACCESS_KEY_ID": "STATICKEY", - "AWS_SECRET_ACCESS_KEY": "STATICSECRET", - "AWS_REGION_NAME": "us-east-1", - } - if session_token: - aws_attrs["AWS_SESSION_TOKEN"] = session_token - settings.aws = type("AWS", (), aws_attrs)() - return settings + overrides["aws.AWS_SESSION_TOKEN"] = session_token + return _base_settings(overrides) -def _frozen_creds( - access_key="FAKE-KEY", - secret_key="FAKE-SECRET", - token=None, -): +def _frozen_creds(access_key="IMDS-KEY", secret_key="IMDS-SECRET", token=None): frozen = MagicMock() frozen.access_key = access_key frozen.secret_key = secret_key @@ -96,586 +56,2026 @@ def _frozen_creds( return frozen -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - -@pytest.fixture(autouse=True) -def clean_aws_env(monkeypatch): - """Ensure AWS env vars don't bleed between tests.""" - for var in ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", - "AWS_SESSION_TOKEN", "AWS_REGION_NAME", "AWS_USE_IMDS", - "AWS_BEDROCK_RUNTIME_ENDPOINT"): - monkeypatch.delenv(var, raising=False) +def _mock_response(): + mock = MagicMock() + mock.__getitem__ = lambda self, key: { + "choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}] + }[key] + mock.dict.return_value = {"choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}]} + return mock @pytest.fixture(autouse=True) -def default_settings(monkeypatch): +def isolate_aws(monkeypatch): + mixin = litellm_handler.BedrockMantleAuthMixin + saved_mixin_signer = mixin.sign_request if mixin is not None else None + saved_mixin_token_resolver = mixin._resolve_bearer_token if mixin is not None else None + saved_original_signer = litellm_handler._bedrock_mantle_sign_request + saved_original_token_resolver = litellm_handler._bedrock_mantle_resolve_bearer_token + marker = litellm_handler._BEDROCK_MANTLE_ORIGINAL_SIGNER + resolver_marker = litellm_handler._BEDROCK_MANTLE_ORIGINAL_TOKEN_RESOLVER + missing = object() + saved_marker = getattr(litellm_handler._sign_bedrock_mantle_request, marker, missing) + saved_resolver_marker = getattr( + litellm_handler._resolve_bedrock_mantle_bearer_token, + resolver_marker, + missing, + ) + for variable in ( + *litellm_handler.AWS_CREDENTIAL_CHAIN_ENV_VARS, + "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "AWS_REGION_NAME", "AWS_REGION", + "AWS_DEFAULT_REGION", "AWS_USE_IMDS", "AWS_BEDROCK_RUNTIME_ENDPOINT", "BEDROCK_MANTLE_API_KEY", + "AWS_BEARER_TOKEN_BEDROCK", "BEDROCK_MANTLE_API_BASE", "BEDROCK_MANTLE_REGION", "OPENAI_API_KEY", + ): + monkeypatch.delenv(variable, raising=False) monkeypatch.setattr(litellm_handler, "get_settings", lambda: _base_settings()) + yield + if mixin is not None: + mixin.sign_request = saved_mixin_signer + mixin._resolve_bearer_token = staticmethod(saved_mixin_token_resolver) + litellm_handler._bedrock_mantle_sign_request = saved_original_signer + litellm_handler._bedrock_mantle_resolve_bearer_token = saved_original_token_resolver + if saved_marker is missing: + if hasattr(litellm_handler._sign_bedrock_mantle_request, marker): + delattr(litellm_handler._sign_bedrock_mantle_request, marker) + else: + setattr(litellm_handler._sign_bedrock_mantle_request, marker, saved_marker) + if saved_resolver_marker is missing: + if hasattr(litellm_handler._resolve_bedrock_mantle_bearer_token, resolver_marker): + delattr(litellm_handler._resolve_bedrock_mantle_bearer_token, resolver_marker) + else: + setattr( + litellm_handler._resolve_bedrock_mantle_bearer_token, + resolver_marker, + saved_resolver_marker, + ) -# --------------------------------------------------------------------------- -# __init__ — credential resolution -# --------------------------------------------------------------------------- +@pytest.fixture +def aws_session(): + session = MagicMock() + session.get_credentials.return_value.get_frozen_credentials.return_value = _frozen_creds() + session.region_name = "us-east-1" + with patch("boto3.Session", return_value=session): + yield session -class TestImdsInit: - def test_bedrock_runtime_endpoint_from_settings(self, monkeypatch): - """aws.AWS_BEDROCK_RUNTIME_ENDPOINT in settings is exported to the environment for litellm.""" - endpoint = "https://vpce-0c618f5a6c9f4912f-jsh4i6at.bedrock-runtime.eu-central-1.vpce.amazonaws.com" - settings = _base_settings(extra_get=lambda key: {"aws.AWS_BEDROCK_RUNTIME_ENDPOINT": endpoint}.get(key)) - settings.aws = type("AWS", (), {"AWS_BEDROCK_RUNTIME_ENDPOINT": endpoint})() - monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) +async def _call(handler, model="bedrock/anthropic.claude-sonnet-4-5"): + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: + mock_call.return_value = _mock_response() + await handler.chat_completion(model=model, system="sys", user="usr") + return mock_call.call_args.kwargs - LiteLLMAIHandler() - assert os.environ["AWS_BEDROCK_RUNTIME_ENDPOINT"] == endpoint +@pytest.mark.asyncio +async def test_bedrock_runtime_endpoint_is_request_local(monkeypatch): + endpoint = "https://vpce.example.com" + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _static_aws_settings(overrides={"aws.AWS_BEDROCK_RUNTIME_ENDPOINT": endpoint}), + ) - def test_bedrock_runtime_endpoint_env_var_not_overridden(self, monkeypatch): - """An already-exported AWS_BEDROCK_RUNTIME_ENDPOINT wins over the settings value.""" - monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", "https://env-configured.example.com") - settings_endpoint = "https://settings-configured.example.com" - settings = _base_settings(extra_get=lambda key: {"aws.AWS_BEDROCK_RUNTIME_ENDPOINT": settings_endpoint}.get(key)) - settings.aws = type("AWS", (), {"AWS_BEDROCK_RUNTIME_ENDPOINT": settings_endpoint})() - monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + kwargs = await _call(LiteLLMAIHandler()) - LiteLLMAIHandler() + assert kwargs["aws_bedrock_runtime_endpoint"] == endpoint + assert "AWS_BEDROCK_RUNTIME_ENDPOINT" not in os.environ - assert os.environ["AWS_BEDROCK_RUNTIME_ENDPOINT"] == "https://env-configured.example.com" - def test_imds_creds_written_to_env(self, monkeypatch): - """When AWS_USE_IMDS=true, boto3 creds are placed in os.environ.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = None - - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - - assert os.environ["AWS_ACCESS_KEY_ID"] == frozen.access_key - assert os.environ["AWS_SECRET_ACCESS_KEY"] == frozen.secret_key - assert handler._aws_imds_mode is True - - def test_imds_session_token_set_when_present(self, monkeypatch): - monkeypatch.setenv("AWS_USE_IMDS", "1") - frozen = _frozen_creds(token="session-token-xyz") - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = None - - with patch("boto3.Session", return_value=mock_session): - LiteLLMAIHandler() +@pytest.mark.asyncio +async def test_bedrock_runtime_endpoint_setting_overrides_ambient_env(monkeypatch): + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", "https://environment.example.com") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _static_aws_settings( + overrides={"aws.AWS_BEDROCK_RUNTIME_ENDPOINT": "https://settings.example.com"} + ), + ) - assert os.environ["AWS_SESSION_TOKEN"] == "session-token-xyz" + kwargs = await _call(LiteLLMAIHandler()) - def test_imds_session_token_cleared_when_absent(self, monkeypatch): - """Stale AWS_SESSION_TOKEN from the environment is removed when IMDS creds have no token.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - monkeypatch.setenv("AWS_SESSION_TOKEN", "stale-token") - frozen = _frozen_creds(token=None) - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = None + assert kwargs["aws_bedrock_runtime_endpoint"] == "https://settings.example.com" + assert os.environ["AWS_BEDROCK_RUNTIME_ENDPOINT"] == "https://environment.example.com" - with patch("boto3.Session", return_value=mock_session): - LiteLLMAIHandler() - assert "AWS_SESSION_TOKEN" not in os.environ +@pytest.mark.asyncio +async def test_bedrock_runtime_endpoint_environment_is_request_local(monkeypatch): + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", "https://tenant-a.example.com") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + monkeypatch.setenv("AWS_REGION", "us-east-1") + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", "https://tenant-b.example.com") - def test_imds_region_auto_resolved(self, monkeypatch): - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = "eu-west-1" + kwargs = await _call(handler) - with patch("boto3.Session", return_value=mock_session): - LiteLLMAIHandler() + assert kwargs["aws_bedrock_runtime_endpoint"] == "https://tenant-a.example.com" + assert os.environ["AWS_BEDROCK_RUNTIME_ENDPOINT"] == "https://tenant-b.example.com" - assert os.environ["AWS_REGION_NAME"] == "eu-west-1" - def test_imds_region_not_overwritten_when_already_set(self, monkeypatch): - monkeypatch.setenv("AWS_USE_IMDS", "true") - monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = "eu-west-1" +@pytest.mark.asyncio +@pytest.mark.skipif( + litellm_handler.BedrockMantleAuthMixin is None, + reason="Installed LiteLLM does not provide BedrockMantleAuthMixin", +) +async def test_bedrock_mantle_runtime_endpoint_is_request_local(monkeypatch): + endpoint = "https://vpce.example.com" + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _static_aws_settings(overrides={"aws.AWS_BEDROCK_RUNTIME_ENDPOINT": endpoint}), + ) - with patch("boto3.Session", return_value=mock_session): - LiteLLMAIHandler() + kwargs = await _call(LiteLLMAIHandler(), model="bedrock_mantle/openai.gpt-oss-120b") + + assert kwargs["aws_bedrock_runtime_endpoint"] == endpoint + assert "AWS_BEDROCK_RUNTIME_ENDPOINT" not in os.environ + + +@pytest.mark.asyncio +async def test_static_credentials_are_forwarded_without_mutating_env(monkeypatch): + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings(session_token="STATIC-TOKEN")) + monkeypatch.setenv("AWS_SESSION_TOKEN", "ambient-token") + handler = LiteLLMAIHandler() + + kwargs = await _call(handler) + + assert kwargs["aws_access_key_id"] == "STATIC-KEY" + assert kwargs["aws_secret_access_key"] == "STATIC-SECRET" + assert kwargs["aws_session_token"] == "STATIC-TOKEN" + assert kwargs["aws_region_name"] == "us-east-1" + assert os.environ["AWS_SESSION_TOKEN"] == "ambient-token" + assert "AWS_ACCESS_KEY_ID" not in os.environ + + +@pytest.mark.parametrize("model", ( + "sagemaker/my-endpoint", + pytest.param( + "bedrock_mantle/openai.gpt-oss-120b", + marks=pytest.mark.skipif( + litellm_handler.BedrockMantleAuthMixin is None, + reason="Installed LiteLLM does not provide BedrockMantleAuthMixin", + ), + ), + pytest.param( + "bedrock_mantle/xai.grok-4.3", + marks=pytest.mark.skipif( + litellm_handler.BedrockMantleAuthMixin is None, + reason="Installed LiteLLM does not provide BedrockMantleAuthMixin", + ), + ), +)) +@pytest.mark.asyncio +async def test_static_credentials_are_forwarded_to_additional_aws_providers(monkeypatch, model): + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + monkeypatch.setenv("AWS_SESSION_TOKEN", "ambient-token") + + kwargs = await _call(LiteLLMAIHandler(), model=model) + + assert kwargs["aws_access_key_id"] == "STATIC-KEY" + assert kwargs["aws_secret_access_key"] == "STATIC-SECRET" + assert kwargs["aws_session_token"] == "" + assert kwargs["aws_region_name"] == "us-east-1" + assert os.environ["AWS_SESSION_TOKEN"] == "ambient-token" + + +@pytest.mark.asyncio +@pytest.mark.skipif( + litellm_handler.BedrockMantleAuthMixin is None, + reason="Installed LiteLLM does not provide BedrockMantleAuthMixin", +) +async def test_bedrock_mantle_signer_receives_request_local_credentials(monkeypatch): + from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig + + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://bedrock-mantle.eu-west-1.api.aws") + handler = LiteLLMAIHandler() + aws_signer = MagicMock() + aws_signer._sign_request.return_value = ({}, None) + config = BedrockMantleChatConfig(aws_signer=aws_signer) + + async def capture_call(**kwargs): + config.sign_request( + headers={}, + optional_params={"aws_access_key_id": None, "aws_region_name": None, "unrelated": "value"}, + request_data={}, + api_base=f"{kwargs['api_base']}/v1/chat/completions", + ) + return _mock_response() - assert os.environ["AWS_REGION_NAME"] == "us-west-2" + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=capture_call): + await handler.chat_completion( + model="bedrock_mantle/openai.gpt-oss-120b", + system="sys", + user="usr", + ) - def test_imds_configured_region_exported_to_env(self, monkeypatch): - """aws.AWS_REGION_NAME in settings must be written to env even in IMDS mode.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = "eu-west-1" # would be used if settings region absent + assert aws_signer._sign_request.call_args.kwargs["optional_params"] == { + "aws_access_key_id": "STATIC-KEY", + "aws_secret_access_key": "STATIC-SECRET", + "aws_session_token": "", + "aws_region_name": "eu-west-1", + "unrelated": "value", + } + assert litellm_handler._bedrock_mantle_request_credentials.get() is None + + +@pytest.mark.asyncio +@pytest.mark.skipif( + litellm_handler.BedrockMantleAuthMixin is None, + reason="Installed LiteLLM does not provide BedrockMantleAuthMixin", +) +async def test_concurrent_bedrock_mantle_signers_keep_credentials_isolated(monkeypatch): + from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig + + monkeypatch.setattr(litellm, "api_key", "another-request-key") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings("TENANT-A")) + tenant_a = LiteLLMAIHandler() + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings("TENANT-B")) + tenant_b = LiteLLMAIHandler() + aws_signer = MagicMock() + aws_signer._sign_request.return_value = ({}, None) + config = BedrockMantleChatConfig(aws_signer=aws_signer) + both_started = asyncio.Event() + started = 0 + + async def capture_call(**kwargs): + nonlocal started + started += 1 + if started == 2: + both_started.set() + await asyncio.wait_for(both_started.wait(), timeout=5) + config.sign_request( + headers={}, + optional_params={}, + request_data={}, + api_base="https://bedrock-mantle.us-east-1.api.aws/v1/chat/completions", + api_key=kwargs["api_key"], + ) + return _mock_response() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=capture_call): + await asyncio.gather( + tenant_a.chat_completion( + model="bedrock_mantle/openai.gpt-oss-120b", + system="sys", + user="usr", + ), + tenant_b.chat_completion( + model="bedrock_mantle/openai.gpt-oss-120b", + system="sys", + user="usr", + ), + ) - monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings()) + signer_credentials = { + (call.kwargs["optional_params"]["aws_access_key_id"], call.kwargs["api_key"]) + for call in aws_signer._sign_request.call_args_list + } + assert signer_credentials == {("TENANT-A-KEY", ""), ("TENANT-B-KEY", "")} + assert litellm_handler._bedrock_mantle_request_credentials.get() is None + assert litellm_handler._bedrock_mantle_block_bearer.get() is False + + +@pytest.mark.asyncio +@pytest.mark.skipif( + litellm_handler.BedrockMantleAuthMixin is None, + reason="Installed LiteLLM does not provide BedrockMantleAuthMixin", +) +async def test_custom_bedrock_mantle_provider_uses_request_local_signer(monkeypatch): + from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig + + settings = _static_aws_settings() + settings.litellm.custom_llm_provider = "bedrock_mantle" + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + aws_signer = MagicMock() + aws_signer._sign_request.return_value = ({}, None) + config = BedrockMantleChatConfig(aws_signer=aws_signer) + + async def capture_call(**kwargs): + config.sign_request( + headers={}, + optional_params={}, + request_data={}, + api_base="https://bedrock-mantle.us-east-1.api.aws/v1/chat/completions", + ) + return _mock_response() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=capture_call): + await LiteLLMAIHandler().chat_completion(model="hosted-model", system="sys", user="usr") + + assert aws_signer._sign_request.call_args.kwargs["optional_params"]["aws_access_key_id"] == "STATIC-KEY" + + +@pytest.mark.asyncio +@pytest.mark.skipif( + litellm_handler.BedrockMantleAuthMixin is None, + reason="Installed LiteLLM does not provide BedrockMantleAuthMixin", +) +async def test_bedrock_mantle_blocks_residual_key_without_disabling_sigv4(monkeypatch): + from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig + + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + monkeypatch.setattr(litellm, "api_key", "another-request-key") + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "late-request-token") + aws_signer = MagicMock() + aws_signer._sign_request.return_value = ({}, None) + config = BedrockMantleChatConfig(aws_signer=aws_signer) + + async def capture_call(**kwargs): + assert kwargs["api_key"] == litellm_handler.DUMMY_LITELLM_API_KEY + config.sign_request( + headers={}, + optional_params={}, + request_data={}, + api_base="https://bedrock-mantle.us-east-1.api.aws/v1/chat/completions", + api_key=kwargs["api_key"], + ) + return _mock_response() - with patch("boto3.Session", return_value=mock_session): - LiteLLMAIHandler() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=capture_call): + await handler.chat_completion( + model="bedrock_mantle/openai.gpt-oss-120b", + system="sys", + user="usr", + ) - # settings region (us-east-1) takes precedence over boto3-resolved region (eu-west-1) - assert os.environ["AWS_REGION_NAME"] == "us-east-1" + assert aws_signer._sign_request.call_args.kwargs["api_key"] == "" + assert litellm_handler._bedrock_mantle_request_credentials.get() is None + assert litellm_handler._bedrock_mantle_block_bearer.get() is False + + +@pytest.mark.asyncio +@pytest.mark.skipif( + litellm_handler.BedrockMantleAuthMixin is None, + reason="Installed LiteLLM does not provide BedrockMantleAuthMixin", +) +async def test_bedrock_mantle_real_dispatch_blocks_residual_bearer(monkeypatch): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.setattr(litellm, "api_key", "another-request-key") + handler = LiteLLMAIHandler.__new__(LiteLLMAIHandler) + handler._request_provider_cache = {} + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "late-request-token") + signer_calls = [] + + def stop_before_network(self, *args, **kwargs): + signer_calls.append(kwargs) + raise RuntimeError("stop before network") + + monkeypatch.setattr(BaseAWSLLM, "_sign_request", stop_before_network) + + with pytest.raises(openai.APIConnectionError): + await handler._acompletion( + model="bedrock_mantle/openai.gpt-oss-120b", + api_key=litellm_handler.DUMMY_LITELLM_API_KEY, + aws_access_key_id="request-key", + aws_secret_access_key="request-secret", + aws_session_token="request-token", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://vpce.example.com", + messages=[{"role": "user", "content": "ping"}], + ) - def test_imds_boto3_creds_stored_for_refresh(self, monkeypatch): - """The boto3 credentials object must be stored so refresh avoids re-reading env vars.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_creds = MagicMock() - mock_creds.get_frozen_credentials.return_value = frozen - mock_session = MagicMock() - mock_session.get_credentials.return_value = mock_creds - mock_session.region_name = "us-east-1" + assert len(signer_calls) == 1 + assert signer_calls[0]["api_key"] == "" + assert signer_calls[0]["optional_params"]["aws_access_key_id"] == "request-key" + assert signer_calls[0]["optional_params"]["aws_bedrock_runtime_endpoint"] == "https://vpce.example.com" + assert not set(litellm_handler.BEDROCK_MANTLE_REQUEST_BODY_EXCLUDED_KEYS).intersection( + signer_calls[0]["request_data"] + ) + assert litellm_handler._bedrock_mantle_request_credentials.get() is None + assert litellm_handler._bedrock_mantle_block_bearer.get() is False + + +@pytest.mark.asyncio +async def test_bedrock_sigv4_ignores_residual_litellm_api_key(monkeypatch): + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + monkeypatch.setattr(litellm, "api_key", "another-request-key") + + kwargs = await _call(LiteLLMAIHandler()) + assert "api_key" not in kwargs + assert kwargs["aws_access_key_id"] == "STATIC-KEY" + assert litellm.api_key == "another-request-key" + + +@pytest.fixture +async def native_bedrock_clients(monkeypatch): + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + cache = LLMClientCache() + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", cache) + monkeypatch.setattr(litellm, "modify_params", False) + + def deny_sync_http(*args, **kwargs): + raise AssertionError("Unexpected synchronous HTTP") + + monkeypatch.setattr(httpx.Client, "send", deny_sync_http) + try: + yield + finally: + for client in {id(value): value for value in cache.cache_dict.values()}.values(): + assert isinstance(client, AsyncHTTPHandler) + await client.close() + + +@pytest.mark.parametrize("timing", ("absent", "initial", "late", "executor")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.asyncio +async def test_native_classic_bedrock_generic_key(monkeypatch, native_bedrock_clients, timing, entrypoint): + from litellm.litellm_core_utils import logging_worker + + settings = _base_settings({ + "aws.AWS_ACCESS_KEY_ID": "owned-key", "aws.AWS_SECRET_ACCESS_KEY": "owned-secret", + "aws.AWS_SESSION_TOKEN": "owned-token", "aws.AWS_REGION_NAME": "eu-west-1", + }) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + monkeypatch.setenv("LITELLM_RUST", "false") + for name in ("callbacks", "input_callback", "success_callback", "failure_callback", + "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, []) + requests_seen = [] + + class TransportReached(BaseException): + pass + + async def send(client, request, **kwargs): + requests_seen.append(request) + assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256 ") + assert "Credential=owned-key/" in request.headers["Authorization"] + assert "/eu-west-1/bedrock/aws4_request" in request.headers["Authorization"] + assert request.headers["X-Amz-Security-Token"] == "owned-token" + assert request.url.host == "bedrock-runtime.eu-west-1.amazonaws.com" + assert "foreign" not in str(request.headers) + raise TransportReached + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + model = "bedrock/anthropic.claude-3-haiku-20240307-v1:0" + try: + # Native completion does not consume the unrelated generic key. + monkeypatch.setattr(litellm, "api_key", "foreign-global-key") + with pytest.raises(TransportReached): + await litellm.acompletion( + model=model, messages=[{"role": "user", "content": "usr"}], + aws_access_key_id="owned-key", aws_secret_access_key="owned-secret", + aws_session_token="owned-token", aws_region_name="eu-west-1", + ) + monkeypatch.setattr(litellm, "api_key", "foreign-global-key" if timing == "initial" else None) + handler = LiteLLMAIHandler() + if timing == "late": + monkeypatch.setattr(litellm, "api_key", "foreign-global-key") + if timing == "executor": + loop = asyncio.get_running_loop() + native_executor = loop.run_in_executor + + def dispatch(executor, function, *args): + monkeypatch.setattr(litellm, "api_key", "foreign-global-key") + return native_executor(executor, function, *args) + + monkeypatch.setattr(loop, "run_in_executor", dispatch) + if entrypoint == "chat": + with pytest.raises(TransportReached): + await handler.chat_completion(model, "sys", "usr") + assert len(requests_seen) == 2 + else: + # The retained system-only probe has a separate native message limit. + with pytest.raises(litellm.BadRequestError, match="bedrock requires at least one non-system message"): + await handler.probe_completion(model) + assert len(requests_seen) == 1 + finally: + await asyncio.sleep(0) + try: + await worker.flush() + finally: + await worker.stop() + + +def test_bedrock_mantle_signer_bridge_forwards_future_arguments(monkeypatch): + forwarded = {} + + def sign_request(self, headers, optional_params, request_data, api_base, **kwargs): + forwarded["self"] = self + forwarded["args"] = (headers, optional_params, request_data, api_base) + forwarded["kwargs"] = kwargs + return "signed" + + monkeypatch.setattr(litellm_handler, "_bedrock_mantle_sign_request", sign_request) + request_credentials = { + "aws_access_key_id": "request-key", + "aws_secret_access_key": "request-secret", + } + token = litellm_handler._bedrock_mantle_request_credentials.set(request_credentials) + signer = object() + try: + result = litellm_handler._sign_bedrock_mantle_request( + signer, + {}, + {"unrelated": "value"}, + {}, + "https://bedrock-mantle.us-east-1.api.aws/v1/chat/completions", + future_argument="future-value", + ) + finally: + litellm_handler._bedrock_mantle_request_credentials.reset(token) + + assert result == "signed" + assert forwarded["self"] is signer + assert forwarded["args"][1] == {"unrelated": "value", **request_credentials} + assert forwarded["kwargs"] == {"future_argument": "future-value"} + + +def test_bedrock_mantle_signer_bridge_fails_closed_without_optional_params(monkeypatch): + sign_request = MagicMock() + monkeypatch.setattr(litellm_handler, "_bedrock_mantle_sign_request", sign_request) + token = litellm_handler._bedrock_mantle_request_credentials.set({"aws_access_key_id": "request-key"}) + try: + with pytest.raises(RuntimeError, match="did not receive optional_params"): + litellm_handler._sign_bedrock_mantle_request(object(), headers={}, request_data={}) + finally: + litellm_handler._bedrock_mantle_request_credentials.reset(token) + + sign_request.assert_not_called() + + +@pytest.mark.asyncio +async def test_health_probe_refreshes_imds_credentials_once(monkeypatch): + handler = LiteLLMAIHandler() + handler._aws_use_imds = True + handler._aws_imds_mode = True + handler._aws_active_creds = { + "aws_access_key_id": "stale-key", + "aws_secret_access_key": "stale-secret", + "aws_region_name": "us-east-1", + } + + def refresh_credentials(): + handler._aws_active_creds = { + "aws_access_key_id": "refreshed-key", + "aws_secret_access_key": "refreshed-secret", + "aws_region_name": "us-east-1", + } + return True - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() + monkeypatch.setattr(handler, "_refresh_aws_imds_credentials", refresh_credentials) + completion = AsyncMock(return_value=_mock_response()) - assert handler._aws_boto3_creds is mock_creds + await handler.probe_completion("bedrock/anthropic.claude-3-sonnet", _completion=completion) - def test_imds_no_creds_from_boto3(self, monkeypatch): - """When boto3 returns no credentials, _aws_imds_mode remains False.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - mock_session = MagicMock() - mock_session.get_credentials.return_value = None - mock_session.region_name = None + completion.assert_awaited_once() + assert completion.call_args.kwargs["aws_access_key_id"] == "refreshed-key" + assert completion.call_args.kwargs["aws_secret_access_key"] == "refreshed-secret" - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - assert handler._aws_imds_mode is False +@pytest.mark.skipif( + litellm_handler.BedrockMantleAuthMixin is None, + reason="Installed LiteLLM does not provide BedrockMantleAuthMixin", +) +def test_bedrock_mantle_signer_bridge_installation_is_idempotent(): + original_signer = litellm_handler._bedrock_mantle_sign_request - def test_imds_boto3_exception_does_not_crash(self, monkeypatch): - """A boto3 exception during credential resolution must not crash __init__.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - mock_session = MagicMock() - mock_session.get_credentials.side_effect = CredentialRetrievalError( - provider="imds", error_msg="connection timeout" + litellm_handler._install_bedrock_mantle_signer_bridge() + litellm_handler._install_bedrock_mantle_signer_bridge() + + assert litellm_handler._bedrock_mantle_sign_request is original_signer + + +@pytest.mark.asyncio +async def test_bedrock_mantle_fails_safely_when_signer_bridge_is_unavailable(monkeypatch): + handler = LiteLLMAIHandler() + completion = AsyncMock() + monkeypatch.setattr(litellm_handler, "BedrockMantleAuthMixin", None) + + with pytest.raises(RuntimeError, match="does not provide the Bedrock Mantle signer bridge"): + await handler._acompletion( + _completion=completion, + model="bedrock_mantle/openai.gpt-oss-120b", + aws_access_key_id="request-key", + aws_secret_access_key="request-secret", + aws_session_token="", + aws_region_name="us-east-1", ) - mock_session.region_name = None - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() # must not raise + completion.assert_not_awaited() - assert handler._aws_imds_mode is False - def test_imds_sts_client_error_does_not_crash(self, monkeypatch): - """A ClientError from STS/AssumeRole (IRSA path) must not crash __init__. - ClientError is not a BotoCoreError subclass, so it needs an explicit catch.""" +@pytest.mark.asyncio +async def test_bedrock_mantle_uses_ambient_credentials_without_signer_bridge(monkeypatch): + handler = LiteLLMAIHandler() + completion = AsyncMock(return_value="response") + monkeypatch.setattr(litellm_handler, "BedrockMantleAuthMixin", None) + + response = await handler._acompletion( + _completion=completion, + model="bedrock_mantle/openai.gpt-oss-120b", + ) + + assert response == "response" + completion.assert_awaited_once_with(model="bedrock_mantle/openai.gpt-oss-120b") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("session_token", (None, "tenant-a-token")) +@pytest.mark.parametrize("region_variable", ("AWS_REGION_NAME", "AWS_REGION")) +async def test_native_aws_environment_is_snapshotted_per_handler(monkeypatch, session_token, region_variable): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "tenant-a-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "tenant-a-secret") + monkeypatch.setenv(region_variable, "ap-northeast-1") + if session_token: + monkeypatch.setenv("AWS_SESSION_TOKEN", session_token) + tenant_a = LiteLLMAIHandler() + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "tenant-b-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "tenant-b-secret") + monkeypatch.setenv("AWS_SESSION_TOKEN", "tenant-b-token") + monkeypatch.setenv(region_variable, "us-west-2") + tenant_b = LiteLLMAIHandler() + + tenant_a_kwargs = await _call(tenant_a) + tenant_b_kwargs = await _call(tenant_b) + + assert tenant_a_kwargs["aws_access_key_id"] == "tenant-a-key" + assert tenant_a_kwargs["aws_secret_access_key"] == "tenant-a-secret" + assert tenant_a_kwargs["aws_session_token"] == (session_token or "") + assert tenant_a_kwargs["aws_region_name"] == "ap-northeast-1" + assert tenant_b_kwargs["aws_access_key_id"] == "tenant-b-key" + assert tenant_b_kwargs["aws_secret_access_key"] == "tenant-b-secret" + assert tenant_b_kwargs["aws_session_token"] == "tenant-b-token" + assert tenant_b_kwargs["aws_region_name"] == "us-west-2" + assert os.environ["AWS_ACCESS_KEY_ID"] == "tenant-b-key" + assert os.environ["AWS_SECRET_ACCESS_KEY"] == "tenant-b-secret" + assert os.environ["AWS_SESSION_TOKEN"] == "tenant-b-token" + assert os.environ[region_variable] == "us-west-2" + + +@pytest.mark.parametrize("use_imds", (False, True)) +@pytest.mark.parametrize("session_token", (None, "", "request-session-token")) +@pytest.mark.parametrize("security_token", (None, "", "request-security-token")) +@pytest.mark.asyncio +async def test_native_aws_environment_token_selection(monkeypatch, use_imds, session_token, security_token): + from botocore.credentials import EnvProvider + from litellm.caching.caching import DualCache + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + environment = { + "AWS_ACCESS_KEY_ID": "request-key", "AWS_SECRET_ACCESS_KEY": "request-secret", + "AWS_REGION": "us-east-1", "AWS_EC2_METADATA_DISABLED": "true", + "AWS_CONFIG_FILE": os.devnull, "AWS_SHARED_CREDENTIALS_FILE": os.devnull, + } + if use_imds: + environment["AWS_USE_IMDS"] = "true" + for name, value in (("AWS_SESSION_TOKEN", session_token), ("AWS_SECURITY_TOKEN", security_token)): + if value is not None: + environment[name] = value + for name, value in environment.items(): + monkeypatch.setenv(name, value) + if use_imds: + # Use only the artificial mapping, never ambient expiration/account data. + native = EnvProvider(environ=environment).load().get_frozen_credentials() + assert native.token == (security_token or session_token or None) + else: + source = BaseAWSLLM() + source.iam_cache = DualCache() + native = source.get_credentials(aws_region_name="us-east-1").get_frozen_credentials() + assert native.token == session_token + + kwargs = await _call(LiteLLMAIHandler()) + + assert kwargs["aws_session_token"] == (native.token or "") + assert kwargs["aws_access_key_id"] == native.access_key + assert kwargs["aws_secret_access_key"] == native.secret_key + + +@pytest.mark.parametrize("model", ("sagemaker_chat/model", "sagemaker_nova/model")) +@pytest.mark.asyncio +async def test_sagemaker_chat_rejects_bedrock_bearer_before_litellm(monkeypatch, model): + from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig + + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + monkeypatch.setenv("AWS_REGION", "us-east-1") + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-only-token") + config = SagemakerChatConfig() + headers, _ = config.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "request-key", + "aws_secret_access_key": "request-secret", + "aws_region_name": "us-east-1", + }, + request_data={}, + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/model/invocations", + api_key=litellm_handler.DUMMY_LITELLM_API_KEY, + ) + assert headers["Authorization"] == "Bearer bedrock-only-token" + + completion = AsyncMock(return_value=_mock_response()) + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new=completion): + with pytest.raises(ValueError, match="Bedrock bearer token fallback for SageMaker"): + await LiteLLMAIHandler().chat_completion( + model=model, + system="sys", + user="usr", + ) + + completion.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_aws_request_without_credential_snapshot_fails_closed(monkeypatch): + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "another-request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "another-request-secret") + + with pytest.raises(ValueError, match="AWS credentials were not resolved for this request"): + await _call(handler) + + +@pytest.mark.asyncio +async def test_native_aws_environment_without_region_fails_closed(monkeypatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_REGION", "another-request-region") + + await _call(handler, model="gpt-4o") + with pytest.raises(ValueError, match="AWS region was not resolved for this request"): + await _call(handler) + + +@pytest.mark.parametrize(("model", "model_id", "expected"), ( + ("bedrock/converse/eu-west-1/model", None, "eu-west-1"), + ("converse/eu-west-1/model", None, "eu-west-1"), + ("bedrock/arn:aws-cn:bedrock:cn-north-1:123:inference-profile/model", None, "cn-north-1"), + ("converse/eu-west-1/model", "arn:aws-us-gov:bedrock:us-gov-west-1:123:inference-profile/model", "us-gov-west-1"), + ("bedrock/eu-west-1/model", "eu-west-2/model", None), + ("bedrock/eu-west-1/model", "arn:aws:bedrock::123:inference-profile/model", None), + ("bedrock/arn:aws:bedrock:EU-WEST-1:123:inference-profile/model", None, None), + ("bedrock/arn:aws:iam:eu-west-1:123:role/model", None, None), + ("bedrock/unknown-region/model", None, None), + ("bedrock/arn:aws:bedrock:eu-west-1.example:123:inference-profile/model", None, None), + ("bedrock/eu-west-1/model", 123, None), + ("bedrock/invoke/eu-west-1/model", None, None), + ("bedrock/unknown/arn:aws:bedrock:eu-west-1:123:imported-model/model", None, None), + ("bedrock/invoke/llama/arn:aws:bedrock:%65u-west-1:123:imported-model/model", None, None), + ("bedrock/https://example.com/arn:aws:bedrock:eu-west-1:123:imported-model/model", None, None), + ("bedrock/eu-west-1/model", "invoke/arn:aws:bedrock:eu-west-1:123:imported-model/model", None), +)) +def test_bedrock_model_region_grammar(model, model_id, expected): + assert litellm_handler._get_bedrock_model_region(model, model_id) == expected + + +@pytest.mark.asyncio +async def test_bedrock_model_region_does_not_mutate_handler_credentials(monkeypatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + handler = LiteLLMAIHandler() + initial = dict(handler._aws_active_creds) + for region in ("eu-west-1", "us-east-1"): + kwargs = await _call(handler, model=f"bedrock/arn:aws:bedrock:{region}:123:inference-profile/model") + assert kwargs["aws_region_name"] == region + assert handler._aws_active_creds == initial + + +@pytest.mark.parametrize("model_source", ("model_id", "model_arn", "region_path")) +@pytest.mark.parametrize("explicit_region", (None, "us-east-1")) +@pytest.mark.parametrize("entrypoint", ("chat", "probe")) +@pytest.mark.asyncio +async def test_bedrock_request_region_from_captured_model(monkeypatch, model_source, explicit_region, entrypoint): + from litellm.litellm_core_utils import logging_worker + + native_model = "anthropic.claude-3-haiku-20240307-v1:0" + arn = f"arn:aws:bedrock:eu-west-1:123456789012:inference-profile/{native_model}" + model = f"bedrock/{native_model}" + settings = {"litellm.model_id": arn} if model_source == "model_id" else {} + if model_source == "model_arn": + model = f"bedrock/{arn}" + elif model_source == "region_path": + model = f"bedrock/eu-west-1/{native_model}" + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _base_settings(settings)) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setenv("LITELLM_RUST", "false") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + if explicit_region: + monkeypatch.setenv("AWS_REGION_NAME", explicit_region) + handler = LiteLLMAIHandler() + settings["litellm.model_id"] = arn.replace("eu-west-1", "ap-south-1") + monkeypatch.setenv("AWS_REGION_NAME", "ap-south-1") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "later-request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "later-request-secret") + expected_region = "eu-west-1" + + if entrypoint == "probe": + # The core retains the old system-only probe; inspect its auth selection + # independently of Bedrock's non-system-message requirement. + completion = AsyncMock(return_value=_mock_response()) + await handler.probe_completion(model, _completion=completion) + kwargs = completion.call_args.kwargs + assert kwargs["aws_region_name"] == expected_region + assert kwargs["aws_access_key_id"] == "request-key" + assert kwargs["aws_secret_access_key"] == "request-secret" + if model_source == "model_id": + assert kwargs["model_id"] == arn + return + + requests_seen = [] + + async def send(client, request, **kwargs): + requests_seen.append(request) + return httpx.Response(200, request=request, json={ + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + "metrics": {"latencyMs": 1}, + }) + + monkeypatch.setattr(httpx.AsyncClient, "send", send) + worker = logging_worker.LoggingWorker() + atexit.unregister(worker._flush_on_exit) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + try: + assert await handler.chat_completion(model, "sys", "usr") == ("ok", "stop") + finally: + await asyncio.sleep(0) + try: + await worker.flush() + finally: + await worker.stop() + assert len(requests_seen) == 1 + request = requests_seen[0] + assert request.url.host == f"bedrock-runtime.{expected_region}.amazonaws.com" + assert f"/{expected_region}/bedrock/aws4_request" in request.headers["Authorization"] + assert "Credential=request-key/" in request.headers["Authorization"] + assert unquote(request.url.path) == f"/model/{native_model if model_source == 'region_path' else arn}/converse" + + +@pytest.mark.parametrize( + ("environment_variable", "settings_region"), + ( + ("AWS_DEFAULT_REGION", None), + (None, "eu-west-1"), + ), +) +@pytest.mark.asyncio +async def test_native_aws_environment_snapshots_supported_region_sources( + monkeypatch, + environment_variable, + settings_region, +): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + if environment_variable: + monkeypatch.setenv(environment_variable, "eu-west-1") + if settings_region: + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _base_settings({"aws.AWS_REGION_NAME": settings_region}), + ) + + kwargs = await _call(LiteLLMAIHandler()) + + assert kwargs["aws_region_name"] == "eu-west-1" + + +@pytest.mark.parametrize( + ("settings_region", "expected_region"), + ( + ("eu-central-1", "eu-central-1"), + (None, "ap-northeast-1"), + ), +) +@pytest.mark.asyncio +async def test_native_aws_environment_region_source_precedence(monkeypatch, settings_region, expected_region): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + monkeypatch.setenv("AWS_REGION", "ap-northeast-1") + monkeypatch.setenv("AWS_DEFAULT_REGION", "us-west-2") + if settings_region: + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _base_settings({"aws.AWS_REGION_NAME": settings_region}), + ) + + kwargs = await _call(LiteLLMAIHandler()) + + assert kwargs["aws_region_name"] == expected_region + + +@pytest.mark.parametrize("missing_variable", ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY")) +@pytest.mark.asyncio +async def test_incomplete_native_aws_environment_credentials_only_raise_for_aws(monkeypatch, missing_variable): + values = { + "AWS_ACCESS_KEY_ID": "ambient-key", + "AWS_SECRET_ACCESS_KEY": "ambient-secret", + } + values.pop(missing_variable) + for variable, value in values.items(): + monkeypatch.setenv(variable, value) + + handler = LiteLLMAIHandler() + + await _call(handler, model="gpt-4o") + with pytest.raises(ValueError, match="AWS environment credentials are incomplete"): + await _call(handler) + + +@pytest.mark.parametrize("use_imds", (False, True)) +@pytest.mark.asyncio +async def test_incomplete_aws_environment_does_not_block_static_credentials(monkeypatch, use_imds): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "incomplete-environment-key") + if use_imds: monkeypatch.setenv("AWS_USE_IMDS", "true") - mock_session = MagicMock() - mock_session.get_credentials.side_effect = ClientError( - error_response={"Error": {"Code": "AccessDenied", "Message": "Not authorized to assume role"}}, - operation_name="AssumeRole", + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + + with patch("boto3.Session") as session_factory: + kwargs = await _call(LiteLLMAIHandler()) + + session_factory.assert_not_called() + assert kwargs["aws_access_key_id"] == "STATIC-KEY" + assert kwargs["aws_secret_access_key"] == "STATIC-SECRET" + assert kwargs["aws_session_token"] == "" + assert kwargs["aws_region_name"] == "us-east-1" + + +@pytest.mark.parametrize("missing_variable", ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY")) +@pytest.mark.asyncio +async def test_bedrock_mantle_bearer_ignores_incomplete_sigv4_environment(monkeypatch, missing_variable): + values = { + "AWS_ACCESS_KEY_ID": "ambient-key", + "AWS_SECRET_ACCESS_KEY": "ambient-secret", + } + values.pop(missing_variable) + for variable, value in values.items(): + monkeypatch.setenv(variable, value) + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "request-bearer-token") + + kwargs = await _call(LiteLLMAIHandler(), model="bedrock_mantle/openai.gpt-oss-120b") + + assert kwargs["api_key"] == "request-bearer-token" + assert "aws_access_key_id" not in kwargs + assert "aws_secret_access_key" not in kwargs + + +@pytest.mark.asyncio +async def test_bedrock_mantle_bearer_ignores_complete_sigv4_environment_without_region(monkeypatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "request-secret") + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "request-bearer-token") + + kwargs = await _call(LiteLLMAIHandler(), model="bedrock_mantle/openai.gpt-oss-120b") + + assert kwargs["api_key"] == "request-bearer-token" + assert "aws_access_key_id" not in kwargs + assert "aws_secret_access_key" not in kwargs + assert kwargs["aws_region_name"] == "us-east-1" + + +@pytest.mark.parametrize("variable", ( + "BEDROCK_MANTLE_REGION", "AWS_REGION_NAME", "AWS_REGION", "AWS_DEFAULT_REGION", +)) +@pytest.mark.asyncio +async def test_imds_mantle_bearer_rejects_uncaptured_region(monkeypatch, aws_session, variable): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "request-bearer-token") + handler = LiteLLMAIHandler() + monkeypatch.setenv(variable, "us-west-2") + + with pytest.raises(ValueError, match="Refusing live aws_region_name environment fallback"): + await _call(handler, model="bedrock_mantle/openai.gpt-oss-120b") + + +@pytest.mark.parametrize("probe", (False, True)) +@pytest.mark.asyncio +async def test_bedrock_mantle_bearer_uses_request_token_without_imds_refresh(monkeypatch, aws_session, probe): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "request-bearer-token") + handler = LiteLLMAIHandler() + refresh = MagicMock(side_effect=AssertionError("SigV4 refresh should not run for bearer authentication")) + monkeypatch.setattr(handler, "_refresh_aws_imds_credentials", refresh) + + if probe: + completion = AsyncMock(return_value=_mock_response()) + await handler.probe_completion( + "bedrock_mantle/openai.gpt-oss-120b", + _completion=completion, ) - mock_session.region_name = None + kwargs = completion.call_args.kwargs + else: + kwargs = await _call(handler, model="bedrock_mantle/openai.gpt-oss-120b") + + aws_session.get_credentials.assert_called_once_with() + refresh.assert_not_called() + assert kwargs["api_key"] == "request-bearer-token" + assert "aws_access_key_id" not in kwargs + assert "aws_secret_access_key" not in kwargs + + +@pytest.mark.parametrize("probe", (False, True)) +@pytest.mark.asyncio +async def test_bedrock_bearer_uses_request_token_without_imds_refresh(monkeypatch, aws_session, probe): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "request-bearer-token") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + handler = LiteLLMAIHandler() + refresh = MagicMock(side_effect=AssertionError("SigV4 refresh should not run for bearer authentication")) + monkeypatch.setattr(handler, "_refresh_aws_imds_credentials", refresh) + + if probe: + completion = AsyncMock(return_value=_mock_response()) + await handler.probe_completion( + "bedrock/model", + _completion=completion, + ) + kwargs = completion.call_args.kwargs + else: + kwargs = await _call(handler, model="bedrock/model") - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() # must not raise + aws_session.get_credentials.assert_called_once_with() + refresh.assert_not_called() + assert kwargs["api_key"] == "request-bearer-token" + assert kwargs["aws_region_name"] == "us-east-1" + assert "aws_access_key_id" not in kwargs + assert "aws_secret_access_key" not in kwargs - assert handler._aws_imds_mode is False - def test_static_session_token_exported_without_imds(self, monkeypatch): - """Non-IMDS static path exports AWS_SESSION_TOKEN when present in settings.""" - settings = _static_aws_settings(session_token="STS-TOKEN-NOIMDS") - monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) +def test_incomplete_static_credentials_raise(monkeypatch): + settings = _base_settings({"aws.AWS_ACCESS_KEY_ID": "incomplete"}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + with pytest.raises(ValueError, match="AWS credentials are incomplete"): LiteLLMAIHandler() - assert os.environ["AWS_SESSION_TOKEN"] == "STS-TOKEN-NOIMDS" - def test_static_session_token_cleared_without_imds(self, monkeypatch): - """Non-IMDS static path clears stale AWS_SESSION_TOKEN when not in settings.""" - monkeypatch.setenv("AWS_SESSION_TOKEN", "stale-token") - settings = _static_aws_settings() # no session_token - monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) +def test_incomplete_static_credentials_warn_when_imds_is_enabled(monkeypatch, aws_session): + monkeypatch.setenv("AWS_USE_IMDS", "true") + settings = _base_settings({"aws.AWS_ACCESS_KEY_ID": "incomplete"}) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings) + logger = MagicMock() + monkeypatch.setattr(litellm_handler, "get_logger", lambda: logger) - LiteLLMAIHandler() + LiteLLMAIHandler() + + logger.warning.assert_called_once_with( + "AWS_USE_IMDS is set but configured static AWS credentials are incomplete; " + "no static fallback is available" + ) + + +@pytest.mark.asyncio +async def test_imds_credentials_are_captured_during_initialization_without_mutating_env(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_SESSION_TOKEN", "ambient-token") + frozen = _frozen_creds(token="imds-token") + session = MagicMock() + session.get_credentials.return_value.get_frozen_credentials.return_value = frozen + session.region_name = "eu-west-1" + + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + session.get_credentials.assert_called_once_with() + assert handler._aws_boto3_creds is session.get_credentials.return_value + await _call(handler) + + assert handler._aws_active_creds == { + "aws_access_key_id": "IMDS-KEY", + "aws_secret_access_key": "IMDS-SECRET", + "aws_session_token": "imds-token", + "aws_region_name": "eu-west-1", + } + assert handler._aws_imds_mode is True + assert os.environ["AWS_SESSION_TOKEN"] == "ambient-token" + assert "AWS_ACCESS_KEY_ID" not in os.environ + + +@pytest.mark.asyncio +async def test_imds_mode_snapshots_explicit_environment_credentials(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "tenant-a-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "tenant-a-secret") + monkeypatch.setenv("AWS_SESSION_TOKEN", "tenant-a-token") + monkeypatch.setenv("AWS_REGION", "ap-northeast-1") + + def create_session(**kwargs): + session = MagicMock() + session.get_credentials.return_value.get_frozen_credentials.return_value = _frozen_creds( + access_key=kwargs["aws_access_key_id"], + secret_key=kwargs["aws_secret_access_key"], + token=kwargs.get("aws_session_token"), + ) + session.region_name = kwargs.get("region_name") + return session - assert "AWS_SESSION_TOKEN" not in os.environ + with patch("boto3.Session", side_effect=create_session) as session_factory: + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "tenant-b-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "tenant-b-secret") + monkeypatch.setenv("AWS_SESSION_TOKEN", "tenant-b-token") + kwargs = await _call(handler) + + session_factory.assert_called_once_with( + aws_access_key_id="tenant-a-key", + aws_secret_access_key="tenant-a-secret", + aws_session_token="tenant-a-token", + region_name="ap-northeast-1", + ) + assert kwargs["aws_access_key_id"] == "tenant-a-key" + assert kwargs["aws_secret_access_key"] == "tenant-a-secret" + assert kwargs["aws_session_token"] == "tenant-a-token" + assert kwargs["aws_region_name"] == "ap-northeast-1" + + +@pytest.mark.asyncio +async def test_imds_mode_rejects_late_environment_credentials(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") - def test_no_imds_when_env_var_absent(self, monkeypatch): - """boto3 must never be imported or called when AWS_USE_IMDS is not set.""" - with patch("boto3.Session") as mock_boto3: + def change_environment_after_capture(): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "another-request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "another-request-secret") + return {} + + monkeypatch.setattr( + LiteLLMAIHandler, + "_snapshot_aws_credential_chain_files", + staticmethod(change_environment_after_capture), + ) + with patch("boto3.Session") as session_factory: + with pytest.raises(ValueError, match="Refusing live AWS credential environment fallback"): LiteLLMAIHandler() - mock_boto3.assert_not_called() + session_factory.assert_not_called() - def test_static_keys_stashed_for_fallback(self, monkeypatch): - """Static keys in config are stashed in _aws_static_creds when IMDS mode is active.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = None - monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings()) +@pytest.mark.asyncio +async def test_imds_mode_does_not_rediscover_late_environment_credentials(monkeypatch, aws_session): + monkeypatch.setenv("AWS_USE_IMDS", "true") + aws_session.get_credentials.return_value = None + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "another-request-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "another-request-secret") - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() + with patch("boto3.Session") as session_factory: + with pytest.raises(ValueError, match="AWS credentials were not resolved for this request"): + await _call(handler) - assert handler._aws_static_creds is not None - assert handler._aws_static_creds["AWS_ACCESS_KEY_ID"] == "STATICKEY" - assert handler._aws_static_creds["AWS_SECRET_ACCESS_KEY"] == "STATICSECRET" - assert handler._aws_static_creds["AWS_REGION_NAME"] == "us-east-1" + session_factory.assert_not_called() - def test_static_session_token_stashed(self, monkeypatch): - """AWS_SESSION_TOKEN from static config is included in _aws_static_creds.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = None - monkeypatch.setattr( - litellm_handler, "get_settings", - lambda: _static_aws_settings(session_token="STATIC-SESSION-TOKEN") - ) +@pytest.mark.parametrize("environment_variable", litellm_handler.AWS_CREDENTIAL_CHAIN_ENV_VARS) +@pytest.mark.parametrize("use_environment_credentials", (False, True)) +@pytest.mark.asyncio +async def test_imds_mode_rejects_late_credential_chain_environment( + monkeypatch, + aws_session, + environment_variable, + use_environment_credentials, +): + monkeypatch.setenv("AWS_USE_IMDS", "true") + if use_environment_credentials: + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "tenant-a-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "tenant-a-secret") + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "another-request-value") + + with patch("boto3.Session") as session_factory: + with pytest.raises(ValueError, match="Refusing changed AWS credential-chain environment"): + await _call(handler) + + session_factory.assert_not_called() + + +@pytest.mark.parametrize( + "environment_variable", + ("AWS_ENDPOINT_URL_BEDROCK_RUNTIME", "AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME"), +) +@pytest.mark.asyncio +async def test_imds_mode_rejects_late_service_endpoint_environment(monkeypatch, aws_session, environment_variable): + monkeypatch.setenv("AWS_USE_IMDS", "true") + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "https://another-request.example.com") + + with patch("boto3.Session") as session_factory: + with pytest.raises(ValueError, match="Refusing changed AWS credential-chain environment"): + await _call(handler) + + session_factory.assert_not_called() + + +@pytest.mark.parametrize( + ("model", "environment_variable"), + ( + ("bedrock/anthropic.claude-sonnet-4-5", "AWS_ENDPOINT_URL"), + ("bedrock/anthropic.claude-sonnet-4-5", "AWS_ENDPOINT_URL_BEDROCK_RUNTIME"), + ("sagemaker/my-endpoint", "AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME"), + ), +) +@pytest.mark.asyncio +async def test_static_mode_rejects_late_aws_request_endpoint_environment( + monkeypatch, + model, + environment_variable, +): + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + handler = LiteLLMAIHandler() + monkeypatch.setenv(environment_variable, "https://another-request.example.com") - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() + with pytest.raises(ValueError, match="Refusing changed AWS request endpoint environment"): + await _call(handler, model=model) - assert handler._aws_static_creds.get("AWS_SESSION_TOKEN") == "STATIC-SESSION-TOKEN" - def test_static_keys_applied_when_imds_returns_no_creds(self, monkeypatch): - """When AWS_USE_IMDS is set but boto3 returns None, static keys are applied to env.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - mock_session = MagicMock() - mock_session.get_credentials.return_value = None - mock_session.region_name = None +@pytest.mark.parametrize("environment_variable", ("AWS_SHARED_CREDENTIALS_FILE", "AWS_CONFIG_FILE")) +@pytest.mark.asyncio +async def test_imds_mode_rejects_changed_credential_chain_file(monkeypatch, aws_session, tmp_path, environment_variable): + credential_chain_file = tmp_path / environment_variable.lower() + credential_chain_file.write_text("request-a", encoding="utf-8") + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv(environment_variable, str(credential_chain_file)) + handler = LiteLLMAIHandler() + credential_chain_file.write_text("request-b", encoding="utf-8") - monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings()) + with patch("boto3.Session") as session_factory: + with pytest.raises(ValueError, match="Refusing changed AWS credential-chain file"): + await _call(handler) - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() + session_factory.assert_not_called() - assert handler._aws_imds_mode is False - assert os.environ["AWS_ACCESS_KEY_ID"] == "STATICKEY" - assert os.environ["AWS_SECRET_ACCESS_KEY"] == "STATICSECRET" - assert os.environ["AWS_REGION_NAME"] == "us-east-1" - def test_imds_failed_path_clears_stale_session_token(self, monkeypatch): - """When IMDS fails and static creds have no token, a stale AWS_SESSION_TOKEN is cleared.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - monkeypatch.setenv("AWS_SESSION_TOKEN", "stale-imds-token") - mock_session = MagicMock() - mock_session.get_credentials.return_value = None - mock_session.region_name = None +@pytest.mark.asyncio +async def test_imds_mode_rejects_late_original_ec2_credential_file(monkeypatch, aws_session, tmp_path): + monkeypatch.setenv("AWS_USE_IMDS", "true") + handler = LiteLLMAIHandler() + monkeypatch.setenv("AWS_CREDENTIAL_FILE", str(tmp_path / "credentials")) - monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings()) + with patch("boto3.Session") as session_factory: + with pytest.raises(ValueError, match="Refusing changed AWS credential-chain environment"): + await _call(handler) - with patch("boto3.Session", return_value=mock_session): + session_factory.assert_not_called() + + +def test_original_ec2_credential_file_content_is_isolated(monkeypatch, aws_session, tmp_path): + credentials_file = tmp_path / "credentials" + credentials_file.write_text("AWSAccessKeyId=request-a", encoding="utf-8") + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_CREDENTIAL_FILE", str(credentials_file)) + handler = LiteLLMAIHandler() + credentials_file.write_text("AWSAccessKeyId=request-b", encoding="utf-8") + + with pytest.raises(ValueError, match="Refusing changed AWS credential-chain file"): + handler._validate_aws_credential_chain_environment() + + +def test_original_ec2_credential_file_rejects_non_files(monkeypatch, tmp_path): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_CREDENTIAL_FILE", str(tmp_path)) + + with pytest.raises(ValueError, match="AWS_CREDENTIAL_FILE must reference a regular file"): + LiteLLMAIHandler() + + +def test_credential_chain_file_paths_include_boto_defaults(monkeypatch, tmp_path): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.delenv("AWS_SHARED_CREDENTIALS_FILE", raising=False) + monkeypatch.delenv("AWS_CONFIG_FILE", raising=False) + monkeypatch.delenv("BOTO_CONFIG", raising=False) + + assert LiteLLMAIHandler._aws_credential_chain_file_paths() == ( + str(tmp_path / ".aws" / "credentials"), + str(tmp_path / ".aws" / "config"), + "/etc/boto.cfg", + str(tmp_path / ".boto"), + ) + + +@pytest.mark.parametrize( + "environment_variable", + ("AWS_SHARED_CREDENTIALS_FILE", "AWS_CONFIG_FILE", "BOTO_CONFIG"), +) +def test_credential_chain_file_paths_preserve_empty_selectors(monkeypatch, tmp_path, environment_variable): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv(environment_variable, "") + + paths = LiteLLMAIHandler._aws_credential_chain_file_paths() + + if environment_variable == "AWS_SHARED_CREDENTIALS_FILE": + assert str(tmp_path / ".aws" / "credentials") not in paths + elif environment_variable == "AWS_CONFIG_FILE": + assert str(tmp_path / ".aws" / "config") not in paths + else: + assert "/etc/boto.cfg" not in paths + assert str(tmp_path / ".boto") not in paths + + +@pytest.mark.parametrize( + "environment_variable", + ("AWS_SHARED_CREDENTIALS_FILE", "AWS_CONFIG_FILE", "BOTO_CONFIG"), +) +def test_credential_chain_file_paths_expand_environment_variables(monkeypatch, tmp_path, environment_variable): + monkeypatch.setenv("CREDENTIAL_HOME", str(tmp_path)) + monkeypatch.setenv(environment_variable, "$CREDENTIAL_HOME/credentials") + + assert str(tmp_path / "credentials") in LiteLLMAIHandler._aws_credential_chain_file_paths() + + +def test_original_ec2_credential_file_path_does_not_expand_environment_variables(monkeypatch, tmp_path): + monkeypatch.setenv("CREDENTIAL_HOME", str(tmp_path)) + monkeypatch.setenv("AWS_CREDENTIAL_FILE", "$CREDENTIAL_HOME/credentials") + + paths = LiteLLMAIHandler._aws_credential_chain_file_paths() + + assert os.path.abspath("$CREDENTIAL_HOME/credentials") in paths + assert str(tmp_path / "credentials") not in paths + + +def test_credential_chain_file_fingerprint_skips_non_files(tmp_path): + with patch("builtins.open") as credential_file: + assert LiteLLMAIHandler._fingerprint_aws_credential_chain_file(str(tmp_path)) == ("nonfile",) + + credential_file.assert_not_called() + + +@pytest.mark.asyncio +async def test_imds_mode_rechecks_credential_chain_file_during_initial_resolution(monkeypatch, tmp_path): + credentials_file = tmp_path / "credentials" + credentials_file.write_text("request-a", encoding="utf-8") + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(credentials_file)) + credentials = MagicMock() + credentials.get_frozen_credentials.return_value = _frozen_creds() + session = MagicMock() + + def change_credentials_file(): + credentials_file.write_text("request-b", encoding="utf-8") + return credentials + + session.get_credentials.side_effect = change_credentials_file + session.region_name = "us-east-1" + + with patch("boto3.Session", return_value=session): + with pytest.raises(ValueError, match="Refusing changed AWS credential-chain file"): LiteLLMAIHandler() - assert "AWS_SESSION_TOKEN" not in os.environ - def test_static_keys_applied_when_boto3_raises(self, monkeypatch): - """When AWS_USE_IMDS is set but boto3 throws, static keys are applied to env.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - mock_session = MagicMock() - mock_session.get_credentials.side_effect = CredentialRetrievalError( - provider="imds", error_msg="connection timeout" - ) - mock_session.region_name = None +@pytest.mark.asyncio +async def test_imds_mode_rechecks_credential_chain_environment_before_refresh(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_PROFILE", "tenant-a") + credentials = MagicMock() + credentials.get_frozen_credentials.return_value = _frozen_creds() + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "us-east-1" + + with patch("boto3.Session", return_value=session) as session_factory: + handler = LiteLLMAIHandler() + await _call(handler) + monkeypatch.setenv("AWS_PROFILE", "tenant-b") + with pytest.raises(ValueError, match="Refusing changed AWS credential-chain environment"): + await _call(handler) + + session_factory.assert_called_once_with() + + +@pytest.mark.parametrize( + "profiles", + ( + {"tenant": {"credential_process": "get-credentials"}}, + { + "tenant": {"role_arn": "arn:aws:iam::123456789012:role/tenant", "source_profile": "source"}, + "source": {"credential_process": "get-credentials"}, + }, + ), +) +@pytest.mark.asyncio +async def test_imds_mode_rejects_credential_process_profiles(monkeypatch, profiles): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_PROFILE", "tenant") + session = MagicMock() + session.profile_name = "tenant" + session._session.full_config = {"profiles": profiles} + + with patch("boto3.Session", return_value=session): + with pytest.raises(ValueError, match="credential_process"): + LiteLLMAIHandler() + + session.get_credentials.assert_not_called() + + +@pytest.mark.asyncio +async def test_imds_mode_uses_static_fallback_for_credential_process_profile(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_PROFILE", "tenant") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + session = MagicMock() + session.profile_name = "tenant" + session._session.full_config = { + "profiles": {"tenant": {"credential_process": "get-credentials"}}, + } + + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + await _call(handler) + + session.get_credentials.assert_not_called() + assert handler._aws_imds_mode is False + assert handler._aws_imds_fell_back is True + assert handler._aws_active_creds["aws_access_key_id"] == "STATIC-KEY" + + +@pytest.mark.asyncio +async def test_imds_mode_rejects_credential_chain_change_during_initial_resolution(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + credentials = MagicMock() + credentials.get_frozen_credentials.return_value = _frozen_creds() + session = MagicMock() + + def change_environment(): + monkeypatch.setenv("AWS_PROFILE", "another-request") + return credentials + + session.get_credentials.side_effect = change_environment + session.region_name = "us-east-1" + + with patch("boto3.Session", return_value=session): + with pytest.raises(ValueError, match="Refusing changed AWS credential-chain environment"): + LiteLLMAIHandler() + + +@pytest.mark.asyncio +async def test_imds_mode_rejects_credential_chain_change_during_refresh(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + credentials = MagicMock() + credentials.get_frozen_credentials.return_value = _frozen_creds() + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "us-east-1" + + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + await _call(handler) + + def change_environment(): + monkeypatch.setenv("AWS_PROFILE", "another-request") + return _frozen_creds("ANOTHER-KEY", "ANOTHER-SECRET") + + credentials.get_frozen_credentials.side_effect = change_environment + with pytest.raises(ValueError, match="Refusing changed AWS credential-chain environment"): + await _call(handler) + + +@pytest.mark.asyncio +async def test_ambient_region_precedes_configured_and_session_region(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + session = MagicMock() + session.get_credentials.return_value.get_frozen_credentials.return_value = _frozen_creds() + session.region_name = "eu-west-1" - monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings()) + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + await _call(handler) + + assert handler._aws_active_creds["aws_region_name"] == "us-west-2" + assert os.environ["AWS_REGION_NAME"] == "us-west-2" - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - assert handler._aws_imds_mode is False - assert os.environ["AWS_ACCESS_KEY_ID"] == "STATICKEY" - assert os.environ["AWS_SECRET_ACCESS_KEY"] == "STATICSECRET" +def test_no_imds_lookup_when_disabled(): + with patch("boto3.Session") as session: + LiteLLMAIHandler() + + session.assert_not_called() + + +@pytest.mark.parametrize("error", ( + CredentialRetrievalError(provider="imds", error_msg="unreachable"), + ClientError( + error_response={"Error": {"Code": "AccessDenied", "Message": "denied"}}, + operation_name="AssumeRole", + ), +)) +@pytest.mark.asyncio +async def test_imds_resolution_failure_uses_request_local_static_fallback(monkeypatch, error): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + session = MagicMock() + session.get_credentials.side_effect = error + session.region_name = None + + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + await _call(handler) + + assert handler._aws_imds_mode is False + assert handler._aws_active_creds["aws_access_key_id"] == "STATIC-KEY" + assert handler._aws_active_creds["aws_session_token"] == "" + assert "AWS_ACCESS_KEY_ID" not in os.environ + + +@pytest.mark.parametrize( + ("model", "environment_variable"), + ( + ("bedrock/anthropic.claude-sonnet-4-5", "AWS_ENDPOINT_URL_BEDROCK_RUNTIME"), + ("sagemaker/my-endpoint", "AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME"), + ), +) +@pytest.mark.asyncio +async def test_static_fallback_rejects_late_aws_request_endpoint_environment( + monkeypatch, + model, + environment_variable, +): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + session = MagicMock() + session.get_credentials.side_effect = CredentialRetrievalError(provider="imds", error_msg="unreachable") + session.region_name = None + + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + await _call(handler, model=model) + monkeypatch.setenv(environment_variable, "https://another-request.example.com") + with pytest.raises(ValueError, match="Refusing changed AWS request endpoint environment"): + await _call(handler, model=model) -# --------------------------------------------------------------------------- -# chat_completion — credential refresh and fallback -# --------------------------------------------------------------------------- -class TestImdsCallBehavior: +@pytest.mark.asyncio +async def test_imds_session_creation_failure_uses_request_local_static_fallback(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) - @pytest.mark.parametrize( - "model", - [ - "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - "bedrock_mantle/xai.grok-4.3", - ], + with patch("boto3.Session", side_effect=ProfileNotFound(profile="missing")): + handler = LiteLLMAIHandler() + kwargs = await _call(handler) + + assert handler._aws_imds_mode is False + assert kwargs["aws_access_key_id"] == "STATIC-KEY" + assert kwargs["aws_secret_access_key"] == "STATIC-SECRET" + assert kwargs["aws_region_name"] == "us-east-1" + assert kwargs["aws_session_token"] == "" + assert handler._aws_imds_fell_back is True + assert "AWS_ACCESS_KEY_ID" not in os.environ + + +@pytest.mark.asyncio +async def test_imds_credential_snapshot_oserror_uses_request_local_static_fallback(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + session = MagicMock() + session.get_credentials.return_value.get_frozen_credentials.side_effect = OSError("token file unavailable") + session.region_name = "us-east-1" + + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + kwargs = await _call(handler) + + assert handler._aws_imds_mode is False + assert kwargs["aws_access_key_id"] == "STATIC-KEY" + assert kwargs["aws_secret_access_key"] == "STATIC-SECRET" + assert kwargs["aws_region_name"] == "us-east-1" + assert kwargs["aws_session_token"] == "" + assert handler._aws_imds_fell_back is True + assert "AWS_ACCESS_KEY_ID" not in os.environ + + +@pytest.mark.asyncio +async def test_imds_refresh_updates_request_credentials(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + credentials = MagicMock() + credentials.get_frozen_credentials.side_effect = ( + _frozen_creds(), + _frozen_creds("ROTATED-KEY", "ROTATED-SECRET", "ROTATED-TOKEN"), ) - @pytest.mark.asyncio - async def test_refresh_called_before_bedrock_call(self, monkeypatch, model): - """_refresh_aws_imds_credentials is called before each Bedrock provider call.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = "us-east-1" - - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - - with patch.object(handler, "_refresh_aws_imds_credentials") as mock_refresh, \ - patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_acompletion_response() - await handler.chat_completion( - model=model, - system="sys", user="usr" - ) + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "us-east-1" - mock_refresh.assert_called_once() - assert mock_call.await_args.kwargs["model"] == model + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + await _call(handler) - def test_refresh_uses_stored_creds_not_new_session(self, monkeypatch): - """_refresh_aws_imds_credentials must call get_frozen_credentials on the stored object, - not create a new boto3.Session (which would re-read env vars and return stale creds).""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen1 = _frozen_creds(access_key="FIRST-KEY", secret_key="FIRST-SECRET") - frozen2 = _frozen_creds(access_key="ROTATED-KEY", secret_key="ROTATED-SECRET") - mock_creds = MagicMock() - mock_creds.get_frozen_credentials.side_effect = [frozen1, frozen2] - mock_session = MagicMock() - mock_session.get_credentials.return_value = mock_creds - mock_session.region_name = "us-east-1" - - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - - # boto3.Session should only be called once (in __init__), not during refresh - with patch("boto3.Session") as mock_boto3_refresh: - handler._refresh_aws_imds_credentials() - - mock_boto3_refresh.assert_not_called() - assert os.environ["AWS_ACCESS_KEY_ID"] == "ROTATED-KEY" - assert os.environ["AWS_SECRET_ACCESS_KEY"] == "ROTATED-SECRET" - - def test_refresh_returns_false_and_warns_when_no_stored_creds(self, monkeypatch): - """_refresh_aws_imds_credentials returns False and logs a warning when _aws_boto3_creds is None.""" + assert handler._aws_active_creds["aws_access_key_id"] == "ROTATED-KEY" + assert handler._aws_active_creds["aws_session_token"] == "ROTATED-TOKEN" + assert "AWS_ACCESS_KEY_ID" not in os.environ + + +@pytest.mark.asyncio +async def test_bedrock_call_refreshes_and_forwards_imds_credentials(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + credentials = MagicMock() + credentials.get_frozen_credentials.side_effect = ( + _frozen_creds(), + _frozen_creds("ROTATED-KEY", "ROTATED-SECRET"), + ) + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "us-east-1" + + with patch("boto3.Session", return_value=session): handler = LiteLLMAIHandler() - assert handler._aws_boto3_creds is None - result = handler._refresh_aws_imds_credentials() - assert result is False + kwargs = await _call(handler) + + assert kwargs["aws_access_key_id"] == "ROTATED-KEY" + assert kwargs["aws_secret_access_key"] == "ROTATED-SECRET" + assert kwargs["aws_session_token"] == "" + + +@pytest.mark.asyncio +async def test_imds_call_lock_serializes_network_calls(monkeypatch, aws_session): + monkeypatch.setenv("AWS_USE_IMDS", "true") + handler = LiteLLMAIHandler() + active_calls = 0 + completed_calls = 0 + + async def capture_call(**kwargs): + nonlocal active_calls, completed_calls + assert handler._aws_bedrock_lock.locked() + active_calls += 1 + assert active_calls == 1 + assert kwargs["aws_access_key_id"] == "IMDS-KEY" + await asyncio.sleep(0) + active_calls -= 1 + completed_calls += 1 + return _mock_response() + + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=capture_call): + await asyncio.gather( + handler.chat_completion(model="bedrock/anthropic.claude-sonnet-4-5", system="sys", user="one"), + handler.chat_completion(model="bedrock/anthropic.claude-sonnet-4-5", system="sys", user="two"), + ) - def test_refresh_returns_false_on_exception(self, monkeypatch): - """_refresh_aws_imds_credentials returns False when get_frozen_credentials raises.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_creds = MagicMock() - mock_creds.get_frozen_credentials.side_effect = [ - frozen, - CredentialRetrievalError(provider="imds", error_msg="token expired"), - ] - mock_session = MagicMock() - mock_session.get_credentials.return_value = mock_creds - mock_session.region_name = "us-east-1" - - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - - result = handler._refresh_aws_imds_credentials() - assert result is False - - def test_refresh_returns_false_on_sts_client_error(self, monkeypatch): - """_refresh_aws_imds_credentials returns False on ClientError (STS/AssumeRole path). - ClientError is not a BotoCoreError subclass, so it needs an explicit catch.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_creds = MagicMock() - mock_creds.get_frozen_credentials.side_effect = [ - frozen, - ClientError( - error_response={"Error": {"Code": "AccessDenied", "Message": "Token expired"}}, - operation_name="AssumeRoleWithWebIdentity", - ), - ] - mock_session = MagicMock() - mock_session.get_credentials.return_value = mock_creds - mock_session.region_name = "us-east-1" + assert completed_calls == 2 + assert not handler._aws_bedrock_lock.locked() + aws_session.get_credentials.assert_called_once_with() + assert "AWS_ACCESS_KEY_ID" not in os.environ + + +@pytest.mark.parametrize( + "error", + ( + CredentialRetrievalError(provider="imds", error_msg="expired"), + OSError("token file unavailable"), + ), +) +@pytest.mark.asyncio +async def test_refresh_failure_activates_static_fallback_before_call(monkeypatch, error): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + credentials = MagicMock() + credentials.get_frozen_credentials.side_effect = ( + _frozen_creds(), + error, + ) + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "us-east-1" - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + kwargs = await _call(handler) - result = handler._refresh_aws_imds_credentials() - assert result is False + assert kwargs["aws_access_key_id"] == "STATIC-KEY" + assert kwargs["aws_session_token"] == "" + assert handler._aws_imds_fell_back is True - @pytest.mark.asyncio - async def test_static_fallback_activated_on_refresh_failure(self, monkeypatch): - """When refresh fails and static creds are available, fallback activates before the call.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_creds = MagicMock() - mock_creds.get_frozen_credentials.side_effect = [ - frozen, - CredentialRetrievalError(provider="imds", error_msg="IMDS unreachable"), - ] - mock_session = MagicMock() - mock_session.get_credentials.return_value = mock_creds - mock_session.region_name = "us-east-1" - - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - - handler._aws_static_creds = { - "AWS_ACCESS_KEY_ID": "STATICKEY", - "AWS_SECRET_ACCESS_KEY": "STATICSECRET", - "AWS_REGION_NAME": "us-east-1", - } - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_acompletion_response() - await handler.chat_completion( - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - system="sys", user="usr" +@pytest.mark.asyncio +async def test_bedrock_api_error_retries_with_static_request_credentials(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_REGION_NAME", "eu-west-1") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _static_aws_settings(session_token="STATIC-TOKEN"), + ) + credentials = MagicMock() + credentials.get_frozen_credentials.return_value = _frozen_creds(token="IMDS-TOKEN") + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "us-east-1" + calls = [] + + async def fail_then_succeed(**kwargs): + calls.append(kwargs) + if len(calls) == 1: + raise openai.AuthenticationError( + "Bedrock auth failed", + response=httpx.Response(401, request=httpx.Request("POST", "https://bedrock.example")), + body=None, ) + return _mock_response() - assert handler._aws_imds_fell_back is True - assert os.environ["AWS_ACCESS_KEY_ID"] == "STATICKEY" - - @pytest.mark.asyncio - async def test_refresh_not_called_for_non_bedrock_model(self, monkeypatch): - """_refresh_aws_imds_credentials is NOT called when model is not a Bedrock model.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = "us-east-1" - - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - - with patch.object(handler, "_refresh_aws_imds_credentials") as mock_refresh, \ - patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock) as mock_call: - mock_call.return_value = _mock_acompletion_response() - await handler.chat_completion(model="gpt-4o", system="sys", user="usr") - - mock_refresh.assert_not_called() - - @pytest.mark.asyncio - async def test_fallback_to_static_on_bedrock_failure(self, monkeypatch): - """On Bedrock APIError, _activate_static_aws_fallback is called and the - in-lock retry uses static credentials (second acompletion call).""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = "us-east-1" - - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - - # Manually inject static creds as fallback - handler._aws_static_creds = { - "AWS_ACCESS_KEY_ID": "STATICKEY", - "AWS_SECRET_ACCESS_KEY": "STATICSECRET", - "AWS_REGION_NAME": "us-east-1", - } + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=fail_then_succeed): + await handler.chat_completion(model="bedrock/anthropic.claude-sonnet-4-5", system="sys", user="usr") + + assert calls[0]["aws_access_key_id"] == "IMDS-KEY" + assert calls[0]["aws_secret_access_key"] == "IMDS-SECRET" + assert calls[0]["aws_session_token"] == "IMDS-TOKEN" + assert calls[0]["aws_region_name"] == "eu-west-1" + assert calls[1]["aws_access_key_id"] == "STATIC-KEY" + assert calls[1]["aws_secret_access_key"] == "STATIC-SECRET" + assert calls[1]["aws_session_token"] == "STATIC-TOKEN" + assert calls[1]["aws_region_name"] == "us-east-1" + assert handler._aws_imds_fell_back is True + + +@pytest.mark.parametrize(("model", "model_id"), ( + ("bedrock_mantle/openai.gpt-oss-120b", None), + ("bedrock/arn:aws:bedrock:eu-west-1:123:inference-profile/model", None), + ("bedrock/eu-west-1/anthropic.claude-3-haiku-20240307-v1:0", None), + ("bedrock/anthropic.claude-3-haiku-20240307-v1:0", + "arn:aws:bedrock:eu-west-1:123:inference-profile/model"), +)) +@pytest.mark.asyncio +async def test_bedrock_static_fallback_preserves_request_region(monkeypatch, model, model_id): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_REGION_NAME", "eu-west-1") + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://bedrock-mantle.eu-west-1.api.aws/v1") + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _static_aws_settings( + session_token="STATIC-TOKEN", overrides={"litellm.model_id": model_id} if model_id else {}, + ), + ) + credentials = MagicMock() + credentials.get_frozen_credentials.return_value = _frozen_creds(token="IMDS-TOKEN") + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "eu-west-1" + calls = [] + + async def fail_then_succeed(**kwargs): + calls.append(kwargs) + if len(calls) == 1: + raise openai.AuthenticationError( + "Bedrock auth failed", + response=httpx.Response(401, request=httpx.Request("POST", "https://bedrock.example")), + body=None, + ) + return _mock_response() - call_count = 0 - usage = {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18} + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=fail_then_succeed): + await handler.chat_completion( + model=model, + system="sys", + user="usr", + ) - async def flaky_completion(**kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise openai.APIError("Bedrock auth failed", request=MagicMock(), body=None) - return _mock_acompletion_response(usage) + assert calls[0]["aws_access_key_id"] == "IMDS-KEY" + assert calls[0]["aws_region_name"] == "eu-west-1" + assert calls[1]["aws_access_key_id"] == "STATIC-KEY" + assert calls[1]["aws_region_name"] == "eu-west-1" + assert handler._aws_imds_fell_back is True + + +@pytest.mark.parametrize( + "error", + ( + openai.AuthenticationError( + "Bedrock auth failed", + response=httpx.Response(401, request=httpx.Request("POST", "https://bedrock.example")), + body=None, + ), + openai.APIConnectionError( + message="Bedrock - ExpiredTokenException: The security token has expired", + request=httpx.Request("POST", "https://bedrock.example"), + ), + openai.RateLimitError( + "rate limited", + response=httpx.Response(429, request=httpx.Request("POST", "https://bedrock.example")), + body=None, + ), + openai.APITimeoutError(request=httpx.Request("POST", "https://bedrock.example")), + openai.PermissionDeniedError( + "Bedrock - AccessDeniedException: User is not authorized to perform this action", + response=httpx.Response(403, request=httpx.Request("POST", "https://bedrock.example")), + body=None, + ), + ), +) +@pytest.mark.asyncio +async def test_health_probe_does_not_retry_api_errors(monkeypatch, error): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + credentials = MagicMock() + credentials.get_frozen_credentials.return_value = _frozen_creds() + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "us-east-1" + completion = AsyncMock(side_effect=error) + + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + with pytest.raises(type(error)): + await handler.probe_completion( + model="bedrock/anthropic.claude-sonnet-4-5", + _completion=completion, + ) - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - side_effect=flaky_completion): + completion.assert_awaited_once() + assert completion.call_args.kwargs["aws_access_key_id"] == "IMDS-KEY" + assert handler._aws_imds_fell_back is False + assert not handler._aws_bedrock_lock.locked() + + +@pytest.mark.parametrize( + "error", + ( + openai.APITimeoutError(request=httpx.Request("POST", "https://bedrock.example")), + openai.PermissionDeniedError( + "Bedrock - AccessDeniedException: User is not authorized to perform this action", + response=httpx.Response(403, request=httpx.Request("POST", "https://bedrock.example")), + body=None, + ), + ), +) +@pytest.mark.asyncio +async def test_chat_completion_preserves_static_fallback_for_api_errors(monkeypatch, error): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + credentials = MagicMock() + credentials.get_frozen_credentials.return_value = _frozen_creds() + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "us-east-1" + completion = AsyncMock(side_effect=(error, _mock_response())) + + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", completion): await handler.chat_completion( - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - system="sys", user="usr" + model="bedrock/anthropic.claude-sonnet-4-5", + system="sys", + user="usr", ) - assert call_count == 2 - assert handler._aws_imds_fell_back is True - assert os.environ["AWS_ACCESS_KEY_ID"] == "STATICKEY" - assert os.environ["AWS_SECRET_ACCESS_KEY"] == "STATICSECRET" + assert completion.await_count == 2 + assert completion.await_args_list[0].kwargs["aws_access_key_id"] == "IMDS-KEY" + assert completion.await_args_list[1].kwargs["aws_access_key_id"] == "STATIC-KEY" + assert handler._aws_imds_fell_back is True + + +@pytest.mark.asyncio +async def test_chat_completion_rate_limit_does_not_switch_aws_credentials(monkeypatch): + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) + credentials = MagicMock() + credentials.get_frozen_credentials.return_value = _frozen_creds() + session = MagicMock() + session.get_credentials.return_value = credentials + session.region_name = "us-east-1" + error = openai.RateLimitError( + "rate limited", + response=httpx.Response(429, request=httpx.Request("POST", "https://bedrock.example")), + body=None, + ) + completion = AsyncMock(side_effect=error) - @pytest.mark.asyncio - async def test_fallback_not_triggered_without_static_creds(self, monkeypatch): - """If no static fallback credentials exist, APIError propagates normally.""" - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds() - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = "us-east-1" - - with patch("boto3.Session", return_value=mock_session): - handler = LiteLLMAIHandler() - - # No static creds stashed - handler._aws_static_creds = None - - with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", - new_callable=AsyncMock, - side_effect=openai.APIError("auth failed", request=MagicMock(), body=None)): - # tenacity exhausts MODEL_RETRIES and re-raises as RetryError - with pytest.raises((openai.APIError, RetryError)): + with patch("boto3.Session", return_value=session): + handler = LiteLLMAIHandler() + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", completion): + with pytest.raises(openai.RateLimitError): await handler.chat_completion( - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - system="sys", user="usr" + model="bedrock/anthropic.claude-sonnet-4-5", + system="sys", + user="usr", ) + completion.assert_awaited_once() + assert completion.call_args.kwargs["aws_access_key_id"] == "IMDS-KEY" + assert handler._aws_imds_fell_back is False -# --------------------------------------------------------------------------- -# _activate_static_aws_fallback — session token handling -# --------------------------------------------------------------------------- -class TestActivateStaticFallback: +@pytest.mark.asyncio +async def test_non_bedrock_call_does_not_receive_aws_credentials(monkeypatch): + monkeypatch.setattr(litellm_handler, "get_settings", _static_aws_settings) - def _make_handler_in_imds_mode(self, monkeypatch): - monkeypatch.setenv("AWS_USE_IMDS", "true") - frozen = _frozen_creds(token="imds-session-token") - mock_session = MagicMock() - mock_session.get_credentials.return_value.get_frozen_credentials.return_value = frozen - mock_session.region_name = "us-east-1" - with patch("boto3.Session", return_value=mock_session): - return LiteLLMAIHandler() - - def test_restores_static_session_token(self, monkeypatch): - """If static creds include a session token, it is restored in env.""" - handler = self._make_handler_in_imds_mode(monkeypatch) - handler._aws_static_creds = { - "AWS_ACCESS_KEY_ID": "SK", - "AWS_SECRET_ACCESS_KEY": "SS", - "AWS_REGION_NAME": "us-east-1", - "AWS_SESSION_TOKEN": "static-sts-token", - } - handler._activate_static_aws_fallback() - - assert os.environ["AWS_SESSION_TOKEN"] == "static-sts-token" - - def test_clears_session_token_when_static_creds_have_none(self, monkeypatch): - """IMDS session token is removed from env when static creds have no token.""" - handler = self._make_handler_in_imds_mode(monkeypatch) - # IMDS token was set during init - assert os.environ.get("AWS_SESSION_TOKEN") == "imds-session-token" - handler._aws_static_creds = { - "AWS_ACCESS_KEY_ID": "SK", - "AWS_SECRET_ACCESS_KEY": "SS", - "AWS_REGION_NAME": "us-east-1", - # no AWS_SESSION_TOKEN - } - handler._activate_static_aws_fallback() + kwargs = await _call(LiteLLMAIHandler(), model="groq/llama-3.3-70b-versatile") + + assert not any(key.startswith("aws_") for key in kwargs) + + +@pytest.mark.asyncio +async def test_concurrent_handlers_keep_static_aws_credentials_isolated(monkeypatch): + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings("TENANT-A")) + tenant_a = LiteLLMAIHandler() + monkeypatch.setattr(litellm_handler, "get_settings", lambda: _static_aws_settings("TENANT-B")) + tenant_b = LiteLLMAIHandler() + calls = [] - assert "AWS_SESSION_TOKEN" not in os.environ + async def capture_call(**kwargs): + await asyncio.sleep(0) + calls.append(kwargs) + return _mock_response() - def test_sets_fell_back_flag(self, monkeypatch): - handler = self._make_handler_in_imds_mode(monkeypatch) - handler._aws_static_creds = { - "AWS_ACCESS_KEY_ID": "SK", - "AWS_SECRET_ACCESS_KEY": "SS", - "AWS_REGION_NAME": "us-east-1", + with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", side_effect=capture_call): + await asyncio.gather( + tenant_a.chat_completion(model="bedrock/anthropic.claude-sonnet-4-5", system="sys", user="usr"), + tenant_b.chat_completion(model="bedrock/anthropic.claude-sonnet-4-5", system="sys", user="usr"), + ) + + access_keys = {call["aws_access_key_id"] for call in calls} + assert access_keys == {"TENANT-A-KEY", "TENANT-B-KEY"} + assert "AWS_ACCESS_KEY_ID" not in os.environ + + +@pytest.mark.parametrize("provider_name", ("assume-role-with-web-identity", "container-role")) +@pytest.mark.parametrize("drift", (None, "initial", "refresh")) +def test_native_workload_provider_keeps_source_and_token_rotation(monkeypatch, tmp_path, provider_name, drift): + import boto3 + + monkeypatch.setenv("AWS_USE_IMDS", "true") + monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1") + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + for variable in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE", "BOTO_CONFIG"): + monkeypatch.setenv(variable, str(tmp_path / variable)) + token_file = tmp_path / "workload-token" + token_file.write_text("original-token") + other_file = tmp_path / "other-token" + other_file.write_text("other-identity-token") + if provider_name == "container-role": + selector = "AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE" + monkeypatch.setenv("AWS_CONTAINER_CREDENTIALS_FULL_URI", "http://localhost/credentials") + # The file must retain precedence over the alternative environment token. + monkeypatch.setenv("AWS_CONTAINER_AUTHORIZATION_TOKEN", "unused-environment-token") + else: + selector = "AWS_WEB_IDENTITY_TOKEN_FILE" + monkeypatch.setenv("AWS_ROLE_ARN", "arn:aws:iam::123456789012:role/workload") + monkeypatch.setenv(selector, str(token_file)) + session = boto3.Session(region_name="us-east-1") + resolver = session._session.get_component("credential_provider") + provider = resolver.get_provider(provider_name) + seen_tokens = [] + + def response_for(token): + seen_tokens.append(token) + return { + "AccessKeyId": "WORKLOAD-KEY", "SecretAccessKey": "WORKLOAD-SECRET", "Token": "session-token", + "Expiration": (datetime.now(UTC) + timedelta(minutes=5)).isoformat(), } - assert handler._aws_imds_fell_back is False - handler._activate_static_aws_fallback() - assert handler._aws_imds_fell_back is True + + if provider_name == "container-role": + sourcer = resolver.get_provider("assume-role")._credential_sourcer + assert sourcer._get_provider("EcsContainer") is provider + fetcher = MagicMock() + fetcher.retrieve_full_uri.side_effect = lambda uri, headers: response_for(headers["Authorization"]) + provider._fetcher = fetcher + else: + client = MagicMock() + + def assume_role(**kwargs): + result = response_for(kwargs["WebIdentityToken"]) + result["SessionToken"] = result.pop("Token") + return {"Credentials": result} + + client.assume_role_with_web_identity.side_effect = assume_role + provider._client_creator = lambda *args, **kwargs: client + + monkeypatch.setattr(boto3, "Session", lambda **kwargs: session) + original_get = session.get_credentials + + def get_with_drift(): + with monkeypatch.context() as changed: + changed.setenv(selector, str(other_file)) + return original_get() + + if drift == "initial": + monkeypatch.setattr(session, "get_credentials", get_with_drift) + handler = LiteLLMAIHandler() + assert handler._aws_imds_mode is True + assert seen_tokens and set(seen_tokens) == {"original-token"} + if provider_name == "assume-role-with-web-identity": + # Native _get_config, not the environment adapter, records the source feature. + assert "CREDENTIALS_ENV_VARS_STS_WEB_ID_TOKEN" in provider._feature_ids + seen_tokens.clear() + + # Replace the file atomically, as a workload token rotation may do. + replacement = tmp_path / "replacement" + replacement.write_text("rotated-token") + replacement.replace(token_file) + native_credentials = handler._aws_boto3_creds + original_freeze = native_credentials.get_frozen_credentials + + def freeze_with_drift(): + with monkeypatch.context() as changed: + changed.setenv(selector, str(other_file)) + return original_freeze() + + if drift == "refresh": + monkeypatch.setattr(native_credentials, "get_frozen_credentials", freeze_with_drift) + assert handler._refresh_aws_imds_credentials() is True + assert seen_tokens == ["rotated-token"] + assert handler._aws_boto3_creds is native_credentials + assert os.environ[selector] == str(token_file) diff --git a/tests/unittest/test_litellm_max_output_tokens.py b/tests/unittest/test_litellm_max_output_tokens.py index ba6fe18a5c..c7636960f9 100644 --- a/tests/unittest/test_litellm_max_output_tokens.py +++ b/tests/unittest/test_litellm_max_output_tokens.py @@ -12,27 +12,28 @@ import pr_agent.algo.ai_handlers.litellm_ai_handler as litellm_handler -# Environment variables that LiteLLMAIHandler.__init__ reads or mutates: the AWS -# credential path (entered when AWS_USE_IMDS is set) writes the AWS_* variables, -# and OPENAI_API_KEY influences the litellm.api_key fallback. +# Clear these credential environment variables before constructing handlers +# to keep the tests independent of deployment-specific AWS or OpenAI credentials. _HANDLER_ENV_VARS = ( + *litellm_handler.AWS_CREDENTIAL_CHAIN_ENV_VARS, "AWS_USE_IMDS", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "AWS_REGION_NAME", + "AWS_BEARER_TOKEN_BEDROCK", "OPENAI_API_KEY", ) @pytest.fixture(autouse=True) def _restore_litellm_globals(): - """LiteLLMAIHandler.__init__ mutates global litellm/openai state and, when - AWS_USE_IMDS is set, os.environ; snapshot and restore both, and drop - AWS_USE_IMDS so the AWS credential path never runs in these tests.""" + """Clear credential environment and the shared LiteLLM key; restore saved state after each test.""" saved = (litellm.api_key, getattr(litellm, "openai_key", None), openai.api_key) saved_env = {name: os.environ.get(name) for name in _HANDLER_ENV_VARS} - os.environ.pop("AWS_USE_IMDS", None) + for name in _HANDLER_ENV_VARS: + os.environ.pop(name, None) + litellm.api_key = None try: yield finally: @@ -46,9 +47,26 @@ def _restore_litellm_globals(): os.environ[name] = value -def _make_settings(config_values=None): +def test_fixture_clears_and_restores_ambient_litellm_key(monkeypatch): + monkeypatch.setattr(litellm, "api_key", "test-ambient-key") + fixture = _restore_litellm_globals.__wrapped__() + next(fixture) + try: + assert litellm.api_key is None + finally: + with pytest.raises(StopIteration): + next(fixture) + assert litellm.api_key == "test-ambient-key" + + +def _make_settings(config_values=None, openrouter=None, custom_llm_provider=""): """Minimal settings whose `config.get(key, ...)` serves the given dict.""" config_values = config_values or {} + settings_values = { + "aws.AWS_ACCESS_KEY_ID": "test-access-key", + "aws.AWS_SECRET_ACCESS_KEY": "test-secret-key", + "aws.AWS_REGION_NAME": "us-east-1", + } class Config: reasoning_effort = None @@ -64,9 +82,12 @@ def get(self, key, default=None): return type("Settings", (), { "config": Config(), "litellm": type("LiteLLM", (), { + "custom_llm_provider": custom_llm_provider, "get": lambda self, key, default=None: default, })(), - "get": lambda self, key, default=None: default, + "get": lambda self, key, default=None: ( + (openrouter or {}) if key == "openrouter" else settings_values.get(key, default) + ), })() @@ -78,22 +99,41 @@ def _mock_response(): return mock -async def _run(monkeypatch, model, config_values): - monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(config_values)) +async def _run( + monkeypatch, + model, + config_values, + openrouter=None, + custom_llm_provider="", + reserve_default=None, +): + monkeypatch.setattr( + litellm_handler, + "get_settings", + lambda: _make_settings(config_values, openrouter, custom_llm_provider), + ) with patch("pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", new_callable=AsyncMock) as mock_call: mock_call.return_value = _mock_response() handler = litellm_handler.LiteLLMAIHandler() + exposed_limit = handler.get_output_token_limit(model) + exposed_reserve = ( + handler.get_output_token_reserve(model, reserve_default) + if reserve_default is not None + else None + ) await handler.chat_completion(model=model, system="sys", user="usr") - return mock_call.call_args[1] + result = (mock_call.call_args[1], exposed_limit) + return (*result, exposed_reserve) if reserve_default is not None else result class TestMaxOutputTokens: @pytest.mark.asyncio async def test_default_sends_no_max_tokens(self, monkeypatch): - kwargs = await _run(monkeypatch, "bedrock/anthropic.claude-sonnet-5-v1:0", {}) + kwargs, exposed_limit = await _run(monkeypatch, "bedrock/anthropic.claude-sonnet-5-v1:0", {}) assert "max_tokens" not in kwargs + assert exposed_limit == 0 @pytest.mark.asyncio @pytest.mark.parametrize("model", [ @@ -101,28 +141,227 @@ async def test_default_sends_no_max_tokens(self, monkeypatch): "gpt-4o", ]) async def test_positive_value_sent_as_max_tokens(self, monkeypatch, model): - kwargs = await _run(monkeypatch, model, {"max_output_tokens": 16000}) + kwargs, exposed_limit = await _run(monkeypatch, model, {"max_output_tokens": 16000}) assert kwargs["max_tokens"] == 16000 + assert exposed_limit == 16000 @pytest.mark.asyncio async def test_extended_thinking_limit_stays_authoritative(self, monkeypatch): - kwargs = await _run(monkeypatch, "claude-3-7-sonnet-20250219", { + kwargs, exposed_limit = await _run(monkeypatch, "claude-3-7-sonnet-20250219", { "max_output_tokens": 16000, "enable_claude_extended_thinking": True, "extended_thinking_budget_tokens": 2048, "extended_thinking_max_output_tokens": 4096, }) assert kwargs["max_tokens"] == 4096 + assert exposed_limit == 4096 assert kwargs["thinking"] == {"type": "enabled", "budget_tokens": 2048} @pytest.mark.asyncio async def test_string_override_is_coerced(self, monkeypatch): # Dynaconf/env overrides can arrive as strings. - kwargs = await _run(monkeypatch, "gpt-4o", {"max_output_tokens": "16000"}) + kwargs, exposed_limit = await _run(monkeypatch, "gpt-4o", {"max_output_tokens": "16000"}) assert kwargs["max_tokens"] == 16000 + assert exposed_limit == 16000 @pytest.mark.asyncio - @pytest.mark.parametrize("value", ["16k", None, 0, -1]) + @pytest.mark.parametrize("value", ["16k", None, 0, -1, float("nan"), float("inf"), "1" * 5_000]) async def test_unset_or_invalid_values_send_nothing(self, monkeypatch, value): - kwargs = await _run(monkeypatch, "gpt-4o", {"max_output_tokens": value}) + kwargs, exposed_limit = await _run(monkeypatch, "gpt-4o", {"max_output_tokens": value}) + assert "max_tokens" not in kwargs + assert exposed_limit == 0 + + @pytest.mark.asyncio + @pytest.mark.parametrize(("value", "expected"), [(400.5, 400), (True, 1), (False, 0)]) + async def test_existing_int_coercion_is_preserved(self, monkeypatch, value, expected): + kwargs, exposed_limit = await _run(monkeypatch, "gpt-4o", {"max_output_tokens": value}) + + assert exposed_limit == expected + if expected: + assert kwargs["max_tokens"] == expected + else: + assert "max_tokens" not in kwargs + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model", "custom_llm_provider"), + [ + ("openrouter/google/gemini-2.5-pro", ""), + ("google/gemini-2.5-pro", "openrouter"), + ], + ) + async def test_openrouter_only_limit_is_exposed_and_sent( + self, monkeypatch, model, custom_llm_provider + ): + kwargs, exposed_limit = await _run( + monkeypatch, + model, + {}, + openrouter={"max_tokens": 16000}, + custom_llm_provider=custom_llm_provider, + ) + + assert exposed_limit == 16000 + assert kwargs["max_tokens"] == 16000 + + @pytest.mark.asyncio + async def test_openrouter_limit_caps_general_limit(self, monkeypatch): + kwargs, exposed_limit = await _run( + monkeypatch, + "openrouter/google/gemini-2.5-pro", + {"max_output_tokens": 16000}, + openrouter={"max_tokens": 4096}, + ) + + assert exposed_limit == 4096 + assert kwargs["max_tokens"] == 4096 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("config_values", "openrouter", "expected_limit", "expected_reserve"), + [ + ({}, {"reasoning_max_tokens": 4096}, 0, 6096), + ({"max_output_tokens": 8000}, {"reasoning_max_tokens": 4096}, 8000, 8000), + ({}, {"reasoning_max_tokens": 4096, "reasoning_effort": "none"}, 0, 2000), + ], + ) + async def test_openrouter_reasoning_expands_only_the_default_reserve( + self, + monkeypatch, + config_values, + openrouter, + expected_limit, + expected_reserve, + ): + kwargs, exposed_limit, exposed_reserve = await _run( + monkeypatch, + "openrouter/google/gemini-2.5-pro", + config_values, + openrouter=openrouter, + reserve_default=2000, + ) + + assert exposed_limit == expected_limit + assert exposed_reserve == expected_reserve + if openrouter.get("reasoning_effort") == "none": + assert kwargs["extra_body"]["reasoning"] == {"enabled": False} + else: + assert kwargs["extra_body"]["reasoning"] == {"max_tokens": 4096} + + @pytest.mark.asyncio + async def test_openrouter_grok_clamped_none_still_reserves_reasoning(self, monkeypatch): + kwargs, exposed_limit, exposed_reserve = await _run( + monkeypatch, + "openrouter/x-ai/grok-4.6", + {}, + openrouter={"reasoning_max_tokens": 8000, "reasoning_effort": "none"}, + reserve_default=2000, + ) + + assert exposed_limit == 0 + assert exposed_reserve == 10000 + assert kwargs["extra_body"]["reasoning"] == {"max_tokens": 8000} + + @pytest.mark.asyncio + async def test_openrouter_astra_uses_one_capped_output_limit(self, monkeypatch): + kwargs, exposed_limit = await _run( + monkeypatch, + "gpt-6-astra", + {"max_output_tokens": 16000}, + openrouter={"max_tokens": 4096}, + custom_llm_provider="openrouter", + ) + + assert exposed_limit == 4096 + assert kwargs["max_completion_tokens"] == 4096 assert "max_tokens" not in kwargs + + @pytest.mark.asyncio + async def test_openrouter_limit_caps_extended_thinking_limit(self, monkeypatch): + model = "openrouter/anthropic/claude-3-7-sonnet-20250219" + kwargs, exposed_limit = await _run( + monkeypatch, + model, + { + "max_output_tokens": 16000, + "enable_claude_extended_thinking": True, + "extended_thinking_budget_tokens": 1024, + "extended_thinking_max_output_tokens": 4096, + "claude_extended_thinking_models_override": [model], + }, + openrouter={"max_tokens": 2048}, + ) + + assert exposed_limit == 2048 + assert kwargs["thinking"] == {"type": "enabled", "budget_tokens": 1024} + assert kwargs["max_tokens"] == 2048 + + @pytest.mark.asyncio + @pytest.mark.parametrize("adaptive_enabled", [True, False]) + async def test_adaptive_only_model_never_uses_legacy_extended_limit( + self, monkeypatch, adaptive_enabled + ): + model = "anthropic/claude-opus-5" + kwargs, exposed_limit = await _run(monkeypatch, model, { + "max_output_tokens": 16000, + "enable_claude_adaptive_thinking": adaptive_enabled, + "enable_claude_extended_thinking": True, + "extended_thinking_budget_tokens": 2048, + "extended_thinking_max_output_tokens": 4096, + "claude_extended_thinking_models_override": [model], + }) + + assert exposed_limit == 16000 + assert kwargs["max_tokens"] == 16000 + if adaptive_enabled: + assert kwargs["thinking"] == {"type": "adaptive"} + else: + assert "thinking" not in kwargs + + @pytest.mark.asyncio + async def test_general_limit_is_live_while_handler_controls_are_snapshots(self, monkeypatch): + model = "openrouter/google/gemini-2.5-pro" + config_values = {"max_output_tokens": 16000} + openrouter = {"max_tokens": 4096} + active_settings = _make_settings(config_values, openrouter) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + handler = litellm_handler.LiteLLMAIHandler() + + config_values["max_output_tokens"] = 2048 + openrouter["max_tokens"] = 1024 + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_call: + mock_call.return_value = _mock_response() + assert handler.get_output_token_limit(model) == 2048 + await handler.chat_completion(model=model, system="sys", user="usr") + + assert mock_call.call_args.kwargs["max_tokens"] == 2048 + + @pytest.mark.asyncio + async def test_extended_thinking_limit_is_snapshotted_for_accessor_and_request(self, monkeypatch): + model = "claude-3-7-sonnet-20250219" + config_values = { + "enable_claude_extended_thinking": True, + "extended_thinking_budget_tokens": 1024, + "extended_thinking_max_output_tokens": 4096, + } + active_settings = _make_settings(config_values) + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + handler = litellm_handler.LiteLLMAIHandler() + + config_values["extended_thinking_budget_tokens"] = 2048 + config_values["extended_thinking_max_output_tokens"] = 8192 + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_call: + mock_call.return_value = _mock_response() + assert handler.get_output_token_limit(model) == 4096 + await handler.chat_completion(model=model, system="sys", user="usr") + + assert mock_call.call_args.kwargs["thinking"] == {"type": "enabled", "budget_tokens": 1024} + assert mock_call.call_args.kwargs["max_tokens"] == 4096 diff --git a/tests/unittest/test_litellm_openrouter_controls.py b/tests/unittest/test_litellm_openrouter_controls.py index 36b74bbcb7..0b15a53aaa 100644 --- a/tests/unittest/test_litellm_openrouter_controls.py +++ b/tests/unittest/test_litellm_openrouter_controls.py @@ -14,6 +14,7 @@ import litellm import openai import pytest +from dynaconf.utils.boxing import DynaBox from litellm.llms.openrouter.chat.transformation import OpenrouterConfig from litellm.utils import get_llm_provider, get_optional_params @@ -131,6 +132,119 @@ async def test_provider_only_wins_over_order(self, monkeypatch): }) assert kwargs["extra_body"]["provider"] == {"only": ["z-ai"]} + @pytest.mark.asyncio + async def test_controls_are_isolated_between_handlers(self, monkeypatch): + first_controls = { + "provider_only": ["z-ai"], + "reasoning_max_tokens": 1024, + "max_tokens": 4096, + "key": "first-secret", + } + active_settings = _make_settings(first_controls, reasoning_effort="low") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + first_handler = litellm_handler.LiteLLMAIHandler() + + first_controls["provider_only"][0] = "mutated" + active_settings = _make_settings(DynaBox({ + "PROVIDER_ONLY": ["novita"], + "REASONING_EFFORT": "high", + "MAX_TOKENS": 2048, + }), reasoning_effort="high") + second_handler = litellm_handler.LiteLLMAIHandler() + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_call: + mock_call.return_value = _mock_response() + await first_handler.chat_completion( + model="openrouter/google/gemini-2.5-pro", system="sys", user="first" + ) + first_kwargs = mock_call.call_args.kwargs + await second_handler.chat_completion( + model="openrouter/google/gemini-2.5-pro", system="sys", user="second" + ) + second_kwargs = mock_call.call_args.kwargs + + assert first_kwargs["extra_body"]["provider"] == {"only": ["z-ai"]} + assert first_kwargs["extra_body"]["reasoning"] == {"max_tokens": 1024} + assert first_kwargs["max_tokens"] == 4096 + assert second_kwargs["extra_body"]["provider"] == {"only": ["novita"]} + assert second_kwargs["extra_body"]["reasoning"] == {"effort": "high"} + assert second_kwargs["max_tokens"] == 2048 + assert set(first_handler._openrouter_controls) == { + "provider_only", + "provider_order", + "allow_fallbacks", + "reasoning_effort", + "reasoning_max_tokens", + "max_tokens", + } + + @pytest.mark.asyncio + async def test_inherited_reasoning_effort_is_isolated_between_handlers(self, monkeypatch): + active_settings = _make_settings(reasoning_effort="low") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + first_handler = litellm_handler.LiteLLMAIHandler() + + active_settings = _make_settings(reasoning_effort="high") + second_handler = litellm_handler.LiteLLMAIHandler() + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_call: + mock_call.return_value = _mock_response() + await first_handler.chat_completion( + model="openrouter/google/gemini-2.5-pro", system="sys", user="first" + ) + first_kwargs = mock_call.call_args.kwargs + await second_handler.chat_completion( + model="openrouter/google/gemini-2.5-pro", system="sys", user="second" + ) + second_kwargs = mock_call.call_args.kwargs + + assert first_kwargs["extra_body"]["reasoning"] == {"effort": "low"} + assert second_kwargs["extra_body"]["reasoning"] == {"effort": "high"} + + @pytest.mark.asyncio + async def test_gpt5_inherited_reasoning_effort_is_isolated_from_later_settings(self, monkeypatch): + active_settings = _make_settings(reasoning_effort="low") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + handler = litellm_handler.LiteLLMAIHandler() + + active_settings = _make_settings(reasoning_effort="high") + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_call: + mock_call.return_value = _mock_response() + await handler.chat_completion( + model="openrouter/openai/gpt-5.1", system="sys", user="usr" + ) + + kwargs = mock_call.call_args.kwargs + assert kwargs["extra_body"]["reasoning"] == {"effort": "low"} + + @pytest.mark.parametrize("model", ("openai/gpt-5.1", "gemini/gemini-2.5-pro")) + @pytest.mark.asyncio + async def test_direct_reasoning_effort_is_isolated_from_later_settings(self, monkeypatch, model): + active_settings = _make_settings(reasoning_effort="low") + monkeypatch.setattr(litellm_handler, "get_settings", lambda: active_settings) + handler = litellm_handler.LiteLLMAIHandler() + + active_settings = _make_settings(reasoning_effort="high") + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as mock_call: + mock_call.return_value = _mock_response() + await handler.chat_completion(model=model, system="sys", user="usr") + + assert mock_call.call_args.kwargs["reasoning_effort"] == "low" + @pytest.mark.asyncio async def test_reasoning_none_disables(self, monkeypatch): kwargs = await _run(monkeypatch, "openrouter/z-ai/glm-5.2", {"reasoning_effort": "none"}) @@ -273,6 +387,68 @@ async def test_global_reasoning_effort_uses_openrouter_body(self, monkeypatch, m assert kwargs["extra_body"]["reasoning"] == {"effort": "low"} assert kwargs["model"] == model + @pytest.mark.asyncio + async def test_custom_provider_raw_model_uses_openrouter_controls(self, monkeypatch): + kwargs = await _run( + monkeypatch, + "google/gemini-2.5-pro", + {"provider_only": ["google"], "max_tokens": 16000}, + reasoning_effort="low", + custom_llm_provider="openrouter", + ) + + assert kwargs["model"] == "google/gemini-2.5-pro" + assert kwargs["custom_llm_provider"] == "openrouter" + assert "reasoning_effort" not in kwargs + assert kwargs["extra_body"] == { + "provider": {"only": ["google"]}, + "reasoning": {"effort": "low"}, + } + assert kwargs["max_tokens"] == 16000 + + @pytest.mark.asyncio + async def test_custom_provider_raw_gpt5_model_uses_only_openrouter_reasoning(self, monkeypatch): + kwargs = await _run( + monkeypatch, + "openai/gpt-5.1", + {"reasoning_max_tokens": 2048}, + custom_llm_provider="openrouter", + ) + + assert kwargs["extra_body"]["reasoning"] == {"max_tokens": 2048} + assert "reasoning_effort" not in kwargs + assert "allowed_openai_params" not in kwargs + assert "temperature" not in kwargs + + @pytest.mark.asyncio + async def test_custom_provider_raw_gpt5_model_inherits_reasoning_effort(self, monkeypatch): + kwargs = await _run( + monkeypatch, + "openai/gpt-5.1", + {}, + reasoning_effort="low", + custom_llm_provider="openrouter", + ) + + assert kwargs["extra_body"]["reasoning"] == {"effort": "low"} + assert "reasoning_effort" not in kwargs + assert "allowed_openai_params" not in kwargs + assert "temperature" not in kwargs + + @pytest.mark.asyncio + async def test_prefixed_openrouter_gpt5_model_uses_only_openrouter_reasoning(self, monkeypatch): + kwargs = await _run( + monkeypatch, + "openrouter/openai/gpt-5.1", + {}, + reasoning_effort="low", + ) + + assert kwargs["extra_body"]["reasoning"] == {"effort": "low"} + assert "reasoning_effort" not in kwargs + assert "allowed_openai_params" not in kwargs + assert "temperature" not in kwargs + @pytest.mark.asyncio async def test_openrouter_effort_overrides_global_effort(self, monkeypatch): kwargs = await _run( @@ -339,7 +515,7 @@ async def test_reasoning_budget_suppresses_global_effort(self, monkeypatch): assert kwargs["extra_body"]["reasoning"] == {"max_tokens": 2048} def test_litellm_requires_openrouter_reasoning_in_extra_body(self): - """Pin the LiteLLM 1.98.0 workaround boundary so upgrades expose when it can be removed.""" + """Pin the LiteLLM compatibility boundary so upgrades expose when it can be removed.""" with pytest.raises(litellm.UnsupportedParamsError): get_optional_params( model="google/gemini-2.5-pro", @@ -373,7 +549,23 @@ async def test_anthropic_reasoning_budget_warns_without_output_headroom(self, mo assert kwargs["extra_body"]["reasoning"] == {"max_tokens": 2048} assert kwargs["max_tokens"] == 1024 assert any( - "must be greater than reasoning_max_tokens" in call.args[0] + "must be greater than the reasoning budget" in call.args[0] + for call in logger.warning.call_args_list + ) + + @pytest.mark.asyncio + async def test_anthropic_reasoning_effort_warns_without_minimum_output_headroom(self, monkeypatch): + logger = MagicMock() + monkeypatch.setattr(litellm_handler, "get_logger", lambda: logger) + kwargs = await _run( + monkeypatch, + "openrouter/anthropic/claude-3.7-sonnet", + {"reasoning_effort": "high", "max_tokens": 1024}, + ) + assert kwargs["extra_body"]["reasoning"] == {"effort": "high"} + assert kwargs["max_tokens"] == 1024 + assert any( + "must be greater than the reasoning budget (1024)" in call.args[0] for call in logger.warning.call_args_list ) diff --git a/tests/unittest/test_litellm_reasoning_effort.py b/tests/unittest/test_litellm_reasoning_effort.py index 0153812936..5ced98a17e 100644 --- a/tests/unittest/test_litellm_reasoning_effort.py +++ b/tests/unittest/test_litellm_reasoning_effort.py @@ -235,8 +235,8 @@ async def test_gpt5_valid_reasoning_effort_xhigh(self, monkeypatch, mock_logger) mock_logger.info.assert_any_call("Using reasoning_effort='xhigh' for GPT-5 model") @pytest.mark.asyncio - async def test_gpt5_valid_reasoning_effort_max(self, monkeypatch, mock_logger): - """Test GPT-5 with valid reasoning_effort='max' from config.""" + async def test_gpt5_reasoning_effort_max_is_mapped_to_xhigh(self, monkeypatch, mock_logger): + """GPT-5 rejects 'max'; the handler must send its top level 'xhigh' instead.""" fake_settings = create_mock_settings("max") monkeypatch.setattr(litellm_handler, "get_settings", lambda: fake_settings) @@ -254,9 +254,9 @@ async def test_gpt5_valid_reasoning_effort_max(self, monkeypatch, mock_logger): ) call_kwargs = mock_completion.call_args[1] - assert call_kwargs["reasoning_effort"] == "max" + assert call_kwargs["reasoning_effort"] == "xhigh" assert "reasoning_effort" in call_kwargs["allowed_openai_params"] - mock_logger.info.assert_any_call("Using reasoning_effort='max' for GPT-5 model") + mock_logger.info.assert_any_call("Using reasoning_effort='xhigh' for GPT-5 model") @pytest.mark.asyncio async def test_gpt5_valid_reasoning_effort_minimal(self, monkeypatch, mock_logger): @@ -1141,7 +1141,7 @@ def test_clamp_grok_reasoning_effort(self, model, configured, expected): ], ) def test_xai_grok_litellm_reasoning_param_support(self, monkeypatch, model, effort, requires_allowlist): - """Pin LiteLLM 1.98 capability gaps so upgrades expose removable workarounds.""" + """Pin LiteLLM capability gaps so upgrades expose removable workarounds.""" monkeypatch.setattr(litellm, "drop_params", False) bundled_model_cost = GetModelCostMap.load_local_model_cost_map() pinned_model_cost = dict(litellm.model_cost) diff --git a/tests/unittest/test_litellm_run_details.py b/tests/unittest/test_litellm_run_details.py index 0d2355dce1..af03720d89 100644 --- a/tests/unittest/test_litellm_run_details.py +++ b/tests/unittest/test_litellm_run_details.py @@ -191,15 +191,24 @@ def _bare_handler(): handler.no_support_temperature_models = [] handler.support_reasoning_models = [] handler.user_message_only_models = [] + handler._claude_thinking_controls = { + "enable_claude_adaptive_thinking": False, + "enable_claude_extended_thinking": False, + "extended_thinking_budget_tokens": 2048, + "extended_thinking_max_output_tokens": 4096, + } + handler._custom_llm_provider = "" handler._aws_imds_mode = False handler._aws_imds_fell_back = False handler._aws_static_creds = None + handler._aws_active_creds = {} handler._aws_bedrock_lock = None return handler def _streaming_handler(): handler = LiteLLMAIHandler.__new__(LiteLLMAIHandler) + handler._sdk_header_defaults = {"organization": None, "project": None, "custom_headers": {}} handler.streaming_required_models = ["streaming-model"] handler.force_streaming_provider = "" handler.force_streaming_api_base_substrings = [] diff --git a/tests/unittest/test_load_yaml_trailing_text.py b/tests/unittest/test_load_yaml_trailing_text.py new file mode 100644 index 0000000000..3a5d0c4dea --- /dev/null +++ b/tests/unittest/test_load_yaml_trailing_text.py @@ -0,0 +1,96 @@ +"""Drop a sign-off the model added after the wrapper's closing fence.""" +from pr_agent.algo.utils import load_yaml + +SIGN_OFF = "\n\nI reviewed the diff and found nothing else worth flagging." + +# The prompts ask for YAML "and nothing else", but 8 of the 12 user prompts end +# with a dangling open ```yaml, so a reply carries either both fences or, more +# often, a closing one only. +WRAPPED = "```yaml\ncode_suggestions: []\n```" +PRIMED = "code_suggestions: []\n```" + + +def test_a_wrapped_mapping_survives_a_sign_off(): + """The snippet fallback cannot recover this: its pattern requires the + closing fence to end the response.""" + assert load_yaml(WRAPPED + SIGN_OFF, first_key="code_suggestions", + last_key="label") == {"code_suggestions": []} + + +def test_a_primed_mapping_survives_a_sign_off(): + """A reply with a closing fence and no opening one cannot be recovered by + any fence pattern, since there is no pair to match.""" + assert load_yaml(PRIMED + SIGN_OFF, first_key="code_suggestions", + last_key="label") == {"code_suggestions": []} + + +def test_a_sign_off_is_not_folded_into_a_block_scalar(): + """Worse than a parse failure: a single block scalar absorbs the fence and + the remark as part of its value, so the answer is published corrupted.""" + wrapped = "```yaml\nresponse: |\n hello\n```" + SIGN_OFF + + assert load_yaml(wrapped)["response"].strip() == "hello" + + +def test_an_answer_following_a_fenced_example_is_not_truncated(): + """Dropping text after that fence would leave a plain scalar, so the + response must be left alone for the fallbacks to handle.""" + assert load_yaml("```\nexample\n```\nresponse: hello") == {} + + +def test_a_fence_inside_a_block_scalar_is_not_treated_as_the_wrapper(): + """Only a fence at content level closes the wrapper.""" + answer = load_yaml("response: |\n Use this:\n ```python\n print(1)\n ```\n") + + assert answer["response"].count("```") == 2 + + +def test_trailing_whitespace_after_the_fence_changes_nothing(): + assert load_yaml("```yaml\nresponse: |\n hello\n``` \n\n")["response"].strip() == "hello" + + +def test_payload_continuing_past_the_fence_still_fails_loudly(): + """A stray fence mid-answer must not yield a partial mapping. Failing keeps the + retry that a truncated answer would quietly replace.""" + assert load_yaml("a: 1\n```\nb: 2") == {} + + +def test_a_list_continuing_past_the_fence_still_fails_loudly(): + """The /improve shape: a suggestion list truncated to its first item would be + published as if it were complete.""" + reply = "code_suggestions:\n - one\n```\n - two" + + assert load_yaml(reply, first_key="code_suggestions", last_key="label") == {} + + +def test_a_nested_mapping_continuing_past_the_fence_still_fails_loudly(): + assert load_yaml("review:\n score: 1\n```\n effort: 2") == {} + + +def test_a_continuation_followed_by_a_sign_off_still_fails_loudly(): + """The tail is a mapping and then prose, so the tail as a whole does not parse. + Only the shape of its first line marks it as more answer rather than a remark.""" + reply = ("code_suggestions:\n - one\n```\nsecurity_concerns: 'No'\nscore: 8\n\n" + "Hope this helps.") + + assert load_yaml(reply, first_key="code_suggestions", last_key="security_concerns") == {} + + +def test_a_crlf_reply_is_repaired(): + """Provider payloads arrive with CRLF, so the fence line ends with \\r.""" + reply = "```yaml\r\nresponse: |\r\n hello\r\n```\r\n\r\nThat is my answer." + + assert load_yaml(reply)["response"].strip() == "hello" + + +def test_a_sign_off_that_reads_as_a_mapping_is_left_alone(): + """The cost of the rule above: nothing tells 'Note: ...' apart from payload, so + such a reply keeps the pre-existing behaviour rather than risking a truncation.""" + assert load_yaml("a: 1\n```\n\nNote: nothing else stood out.") == {} + + +def test_a_second_fenced_block_is_left_to_the_existing_behaviour(): + """Nothing follows the last fence, so there is no sign-off to drop.""" + reply = "```yaml\nresponse: |\n hello\n```\n\n```\nprint(1)\n```" + + assert "hello" in load_yaml(reply)["response"] diff --git a/tests/unittest/test_local_git_provider.py b/tests/unittest/test_local_git_provider.py index 1da3e9cc0a..1689973fb6 100644 --- a/tests/unittest/test_local_git_provider.py +++ b/tests/unittest/test_local_git_provider.py @@ -1,7 +1,7 @@ import git import pytest -from pr_agent.algo.types import EDIT_TYPE +from pr_agent.algo.types import EDIT_TYPE, FilePatchInfo from pr_agent.config_loader import get_settings from pr_agent.git_providers.local_git_provider import LocalGitProvider from pr_agent.tools.pr_code_suggestions import PRCodeSuggestions @@ -190,6 +190,44 @@ def test_publish_code_suggestions_artifact_includes_partial_coverage(tmp_path): assert "1 of 2 analysis chunks failed" in content +@pytest.mark.asyncio +async def test_mixed_suggestions_stay_in_improve_artifact(tmp_path): + improve_path = tmp_path / "improve.md" + review_path = tmp_path / "review.md" + provider = object.__new__(LocalGitProvider) + provider.improve_path = improve_path + provider.review_path = review_path + provider.diff_files = [FilePatchInfo( + base_file="old()\n", + head_file="old()\n", + patch="", + filename="app.py", + )] + tool = PRCodeSuggestions.__new__(PRCodeSuggestions) + tool.git_provider = provider + tool.progress_response = None + + suggestion = { + "label": "maintainability", + "relevant_file": "app.py", + "suggestion_content": "Use the helper.", + "existing_code": "old()", + "improved_code": "new()", + "score": 8, + } + await tool.push_inline_code_suggestions({"code_suggestions": [ + {**suggestion, "relevant_lines_start": 1, "relevant_lines_end": 1}, + {**suggestion, "suggestion_content": "Keep this advice.", + "relevant_lines_start": 40, "relevant_lines_end": 40}, + ]}) + + content = improve_path.read_text(encoding="utf-8") + assert content.index("Use the helper.") < content.index("Keep this advice.") + assert content.count("```suggestion") == 1 + assert "because the anchored range is outside the file" in content + assert not review_path.exists() + + def test_publish_code_suggestions_uses_custom_heading_without_identity(tmp_path): snapshot = snapshot_settings(["pr_code_suggestions.suggestions_heading"]) improve_path = tmp_path / "improve.md" diff --git a/tests/unittest/test_mosaico_health.py b/tests/unittest/test_mosaico_health.py index 21605d03f5..a379b24956 100644 --- a/tests/unittest/test_mosaico_health.py +++ b/tests/unittest/test_mosaico_health.py @@ -63,7 +63,7 @@ async def test_unhealthy_returns_503(self, monkeypatch): def restore_config_model(): """Restore LLM settings exactly, including originally-absent state.""" snapshot = snapshot_settings( - ["CONFIG.MODEL", "LITELLM.CUSTOM_LLM_PROVIDER"] + ["CONFIG.MODEL", "LITELLM.CUSTOM_LLM_PROVIDER", "OPENAI.KEY", "GROQ.KEY"] ) yield get_settings() restore_settings(snapshot) @@ -72,6 +72,26 @@ def restore_config_model(): class TestHealthCheckGate: """Exercise the REAL health_check() (not the monkeypatched stub) to lock in Fix A.""" + @pytest.mark.asyncio + async def test_health_credentials_remain_request_local(self, monkeypatch, restore_config_model): + restore_config_model.set("LITELLM.CUSTOM_LLM_PROVIDER", "") + restore_config_model.set("GROQ.KEY", "groq-request-key") + restore_config_model.set("OPENAI.KEY", "openai-request-key") + monkeypatch.setattr(litellm, "api_key", "unrelated-global-key") + calls = [] + + async def fake_acompletion(**kwargs): + calls.append(kwargs) + return {"choices": [{"message": {"content": "pong"}}]} + + monkeypatch.setattr(litellm, "acompletion", fake_acompletion) + for model in ("groq/gemma2-9b-it", "openai/gpt-4o"): + restore_config_model.set("CONFIG.MODEL", model) + assert await health_check() == "OK" + + assert [call["api_key"] for call in calls] == ["groq-request-key", "openai-request-key"] + assert litellm.api_key == "unrelated-global-key" + @pytest.mark.asyncio async def test_model_without_stop_probes_live_and_returns_ok( self, monkeypatch, restore_config_model @@ -84,7 +104,7 @@ async def fake_acompletion(**kwargs): called.update(kwargs) return {"choices": [{"message": {"content": "pong"}}]} - # health_check does `import litellm` then `await litellm.acompletion(...)`. + # health_check injects litellm.acompletion into the handler's probe. monkeypatch.setattr(litellm, "acompletion", fake_acompletion) result = await health_check() diff --git a/tests/unittest/test_output_models.py b/tests/unittest/test_output_models.py new file mode 100644 index 0000000000..3381698e2b --- /dev/null +++ b/tests/unittest/test_output_models.py @@ -0,0 +1,341 @@ +import re +from pathlib import Path +from types import UnionType +from typing import get_args, get_origin + +import pytest +from pydantic import StrictInt + +from pr_agent.algo.output_models import ( + CodeDocumentation, + CodeSuggestion, + CodeSuggestionFeedback, + ContributionTimeCostEstimate, + DocHeadingsHelper, + DocHelper, + FileDescription, + FileIdxAndPath, + KeyIssuesComponentLink, + Label, + Labels, + PRCodeSuggestions, + PRCodeSuggestionsFeedback, + PRDescription, + PRDescriptionHeaders, + PRFilesWalkthrough, + PRRankResponses, + PRReview, + PRType, + RelevantSection, + Review, + SubPR, + TicketCompliance, + TodoSection, +) +from pr_agent.algo.utils import load_yaml + + +def _review_fixture(): + issue = { + "relevant_file": "src/app.py", + "issue_header": "Possible Bug", + "issue_content": "The error path can lose the original exception.", + "start_line": 10, + "end_line": 12, + } + return { + "review": { + "estimated_effort_to_review_[1-5]": 3, + "risk_level": "medium", + "merge_recommendation": "merge_with_caution", + "review_priority_files": ["src/app.py"], + "contribution_time_cost_estimate": {"best_case": "45m", "average_case": "2h", "worst_case": "5h"}, + "score": 89, + "relevant_tests": "Yes", + "insights_from_user_answers": "The deployment target is Linux.", + "key_issues_to_review": [issue], + "security_concerns": "No", + "todo_sections": [{"relevant_file": "src/app.py", "line_number": 20, "content": "Remove fallback."}], + "can_be_split": [{"relevant_files": ["src/app.py"], "title": "Improve error handling"}], + "ticket_compliance_check": [{ + "ticket_url": "#123", + "ticket_requirements": "Handle errors.", + "fully_compliant_requirements": "Handle errors.", + "not_compliant_requirements": "", + "requires_further_human_verification": "", + }], + } + } + + +@pytest.mark.parametrize( + ("model", "payload"), + [ + (SubPR, {"relevant_files": ["src/app.py"], "title": "Improve error handling"}), + (KeyIssuesComponentLink, _review_fixture()["review"]["key_issues_to_review"][0]), + (TodoSection, {"relevant_file": "src/app.py", "line_number": 20, "content": "Remove fallback."}), + (TicketCompliance, _review_fixture()["review"]["ticket_compliance_check"][0]), + (ContributionTimeCostEstimate, {"best_case": "45m", "average_case": "2h", "worst_case": "5h"}), + (Review, _review_fixture()["review"]), + (PRReview, _review_fixture()), + (CodeSuggestion, { + "relevant_file": "src/app.py", "language": "python", "existing_code": "return value", + "suggestion_content": "Handle the missing value.", "improved_code": "return value or default", + "one_sentence_summary": "Handle missing values", "label": "possible bug", + }), + (PRCodeSuggestions, {"code_suggestions": [{ + "relevant_file": "src/app.py", "language": "python", "existing_code": "return value", + "suggestion_content": "Handle the missing value.", "improved_code": "return value or default", + "one_sentence_summary": "Handle missing values", "label": "possible bug", + }]}), + (PRCodeSuggestionsFeedback, {"code_suggestions": [{ + "suggestion_summary": "Handle missing values", "relevant_file": "src/app.py", + "relevant_lines_start": 10, "relevant_lines_end": 10, "suggestion_score": 8, + "why": "The change prevents a runtime failure.", + }]}), + (PRDescription, {"type": ["Bug fix"], "description": "Fix a runtime failure.", "title": "Fix runtime failure", + "changes_diagram": "flowchart LR\nA --> B", "pr_files": [{ + "filename": "src/app.py", "changes_summary": "- Handle failures", + "changes_title": "Handle runtime failures", "label": "bug fix", + }]}), + (PRDescriptionHeaders, {"type": ["Tests"], "description": "Cover the new behavior.", + "title": "Describe the test changes", "changes_diagram": ""}), + (FileDescription, {"filename": "src/app.py", "changes_summary": "- Handle failures", + "changes_title": "Handle runtime failures", "label": "bug fix"}), + (PRFilesWalkthrough, {"pr_files": [{ + "filename": "src/app.py", "changes_summary": "- Handle failures", + "changes_title": "Handle runtime failures", "label": "bug fix", + }]}), + (Labels, {"labels": ["Bug fix", "Tests"]}), + (CodeDocumentation, {"Code Documentation": [{ + "relevant file": "src/app.py", "relevant line": 12, "doc placement": "after", + "documentation": "Document the handler.", + }]}), + (PRRankResponses, {"which_response_was_better": 1, "why": "It is clearer.", "score_response1": 9, "score_response2": 7}), + (DocHelper, {"user_question": "How?", "response": "Use the helper.", "relevant_sections": [{ + "file_name": "docs/guide.md", "relevant_section_header_string": "## Usage", + }], "question_is_relevant": 1}), + (DocHeadingsHelper, {"user_question": "How?", "relevant_files_ranking": [{"idx": 0, "file_name": "docs/guide.md"}]}), + ], +) +def test_output_models_validate_complete_fixtures(model, payload): + model.model_validate(payload) + + +def test_review_alias_accepts_prompt_field_name(): + assert Review.model_validate({"key_issues_to_review": [], "estimated_effort_to_review_[1-5]": 3}).estimated_effort_to_review == 3 + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [], "estimated_effort_to_review": 3}) + + +def test_review_rejects_coercible_numeric_types_and_strips_prompt_literals(): + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [], "score": "89"}) + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [], "score": True}) + + review = Review.model_validate({ + "key_issues_to_review": [], + "risk_level": "low\n", + "merge_recommendation": "safe_to_merge\n", + "relevant_tests": "No\n", + }) + assert review.risk_level == "low" + assert review.merge_recommendation == "safe_to_merge" + assert review.relevant_tests == "No" + + +def test_required_label_and_list_constraints_are_enforced(): + suggestion = { + "relevant_file": "src/app.py", "language": "python", "existing_code": "return value", + "suggestion_content": "Handle the missing value.", "improved_code": "return value or default", + "one_sentence_summary": "Handle missing values", + } + with pytest.raises(ValueError): + CodeSuggestion.model_validate(suggestion) + with pytest.raises(ValueError): + PRRankResponses.model_validate({"which_response_was_better": 3, "why": "No", "score_response1": 1, "score_response2": 1}) + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [], "can_be_split": [{"relevant_files": [], "title": "x"}] * 4}) + with pytest.raises(ValueError): + PRDescription.model_validate({"type": ["Tests"], "title": "x", "pr_files": [{ + "filename": "x", "changes_title": "x", "label": "x", + }] * 21}) + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [], "estimated_effort_to_review_[1-5]": 0}) + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [], "score": 101}) + with pytest.raises(ValueError): + PRCodeSuggestionsFeedback.model_validate({"code_suggestions": [{ + "suggestion_summary": "x", "relevant_file": "x", "relevant_lines_start": 1, + "relevant_lines_end": 1, "suggestion_score": 11, "why": "x", + }]}) + with pytest.raises(ValueError): + PRRankResponses.model_validate({"which_response_was_better": 1, "why": "x", "score_response1": 0, "score_response2": 11}) + with pytest.raises(ValueError): + DocHelper.model_validate({"user_question": "x", "response": "x", "relevant_sections": [], "question_is_relevant": 2}) + with pytest.raises(ValueError): + DocHeadingsHelper.model_validate({"user_question": "x", "relevant_files_ranking": [{"idx": -1, "file_name": "x"}]}) + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [], "risk_level": "critical"}) + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [], "merge_recommendation": "approve"}) + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [], "relevant_tests": "Maybe"}) + with pytest.raises(ValueError): + PRDescription.model_validate({"type": [], "title": "x"}) + with pytest.raises(ValueError): + PRDescriptionHeaders.model_validate({"type": [], "title": "x"}) + with pytest.raises(ValueError): + Review.model_validate({"key_issues_to_review": [{ + "relevant_file": "x", "issue_header": "x", "issue_content": "x", + "start_line": 1, "end_line": 1, "unexpected": "x", + }]}) + + +def _split_type_args(value): + parts, depth, start = [], 0, 0 + for index, character in enumerate(value): + if character == "[": + depth += 1 + elif character == "]": + depth -= 1 + elif character == "," and depth == 0: + parts.append(value[start:index]) + start = index + 1 + parts.append(value[start:]) + return parts + + +def _prompt_type_signature(annotation): + annotation = annotation.strip().replace(" ", "") + if annotation.startswith("Optional["): + annotation = annotation[9:-1] + if annotation.startswith("Union["): + return ("union", tuple(sorted( + (_prompt_type_signature(item) for item in _split_type_args(annotation[6:-1])), key=repr + ))) + if annotation.startswith("List["): + return ("list", _prompt_type_signature(annotation[5:-1])) + if annotation.startswith("Literal["): + values = tuple(value.strip().strip('"').strip("'") for value in annotation[8:-1].split(",")) + return ("literal", values) + if annotation == "Label": + return "str" + if annotation == "relevant_section": + return "RelevantSection" + if annotation == "file_idx_and_path": + return "FileIdxAndPath" + return annotation + + +def _model_type_signature(annotation): + origin = get_origin(annotation) + if str(origin) == "typing.Annotated": + return _model_type_signature(get_args(annotation)[0]) + if annotation is StrictInt: + return "int" + if str(origin) == "typing.Literal": + return ("literal", tuple(str(value) for value in get_args(annotation))) + if origin in (list,): + return ("list", _model_type_signature(get_args(annotation)[0])) + if origin in (UnionType,): + args = get_args(annotation) + elif str(origin) == "typing.Union": + args = get_args(annotation) + else: + return getattr(annotation, "__name__", str(annotation).split(".")[-1]) + members = tuple(sorted( + (_model_type_signature(item) for item in args if item is not type(None)), key=repr + )) + return members[0] if len(members) == 1 else ("union", members) + + +PROMPT_MODELS = { + "pr_reviewer_prompts.toml": {"SubPR": SubPR, "KeyIssuesComponentLink": KeyIssuesComponentLink, + "TodoSection": TodoSection, "TicketCompliance": TicketCompliance, + "ContributionTimeCostEstimate": ContributionTimeCostEstimate, "Review": Review, + "PRReview": PRReview}, + "pr_description_prompts.toml": {"FileDescription": FileDescription, "PRDescription": PRDescription}, + "pr_description_only_description_prompts.toml": {"PRDescriptionHeaders": PRDescriptionHeaders}, + "pr_description_only_files_prompts.toml": {"FileDescription": FileDescription, "PRFilesWalkthrough": PRFilesWalkthrough}, + "pr_custom_labels.toml": {"Labels": Labels}, + "pr_evaluate_prompt_response.toml": {"PRRankResponses": PRRankResponses}, + "pr_help_prompts.toml": {"relevant_section": RelevantSection, "DocHelper": DocHelper}, + "pr_help_docs_prompts.toml": {"relevant_section": RelevantSection, "DocHelper": DocHelper}, + "pr_help_docs_headings_prompts.toml": {"file_idx_and_path": FileIdxAndPath, "DocHeadingsHelper": DocHeadingsHelper}, + "code_suggestions/pr_code_suggestions_prompts.toml": {"CodeSuggestion": CodeSuggestion, "PRCodeSuggestions": PRCodeSuggestions}, + "code_suggestions/pr_code_suggestions_prompts_not_decoupled.toml": {"CodeSuggestion": CodeSuggestion, "PRCodeSuggestions": PRCodeSuggestions}, + "code_suggestions/pr_code_suggestions_reflect_prompts.toml": {"CodeSuggestionFeedback": CodeSuggestionFeedback, "PRCodeSuggestionsFeedback": PRCodeSuggestionsFeedback}, +} + + +def test_prompt_fields_are_present_in_output_models(): + root = Path(__file__).parents[2] / "pr_agent" / "settings" + for relative_path, classes in PROMPT_MODELS.items(): + text = (root / relative_path).read_text(encoding="utf-8") + for class_name, model in classes.items(): + class_start = text.index(f"class {class_name}(BaseModel):") + block_start = class_start + len(f"class {class_name}(BaseModel):") + next_class = text.find("\nclass ", block_start) + separator = text.find("=====", block_start) + block_end = min(value for value in (next_class, separator) if value >= 0) + block = text[block_start:block_end] + declared = set(re.findall(r"^ ([A-Za-z_][A-Za-z0-9_\[\]-]*):", block, re.MULTILINE)) + model_fields = set(model.model_fields) + aliases = {field.alias for field in model.model_fields.values() if field.alias} + assert declared <= model_fields | aliases, f"{relative_path}: {class_name} has unmodelled fields" + for field_name, annotation in re.findall( + r"^ ([A-Za-z_][A-Za-z0-9_\[\]-]*):\s*([^=]+)", block, re.MULTILINE + ): + field = next((value for value in model.model_fields.values() + if value.alias == field_name or value.alias is None and value.validation_alias == field_name), None) + if field is None: + field = model.model_fields.get(field_name) + assert field is not None, f"{relative_path}: {class_name} field {field_name} is missing" + assert _prompt_type_signature(annotation) == _model_type_signature(field.annotation), ( + f"{relative_path}: {class_name}.{field_name} type drift" + ) + + +def test_prompt_enum_contracts_are_preserved(): + assert {member.value for member in PRType} == {"Bug fix", "Tests", "Enhancement", "Documentation", "Other"} + assert {member.value for member in Label} == {"Bug fix", "Tests", "Enhancement", "Documentation", "Other"} + + +def test_add_docs_prompt_matches_output_model(): + prompt = (Path(__file__).parents[2] / "pr_agent" / "settings" / "pr_add_docs.toml").read_text(encoding="utf-8") + assert "Code Documentation:" in prompt + assert "relevant file:" in prompt + assert "relevant line:" in prompt + assert "doc placement:" in prompt + assert "documentation:" in prompt + CodeDocumentation.model_validate({"Code Documentation": [{ + "relevant file": "src/app.py", "relevant line": 12, "doc placement": "after", + "documentation": "Document the handler.", + }]}) + + +def test_duplicate_doc_items_are_rejected(): + item = { + "relevant file": "src/app.py", "relevant line": 12, "doc placement": "after", + "documentation": "Document the handler.", + } + with pytest.raises(ValueError): + CodeDocumentation.model_validate({"Code Documentation": [item, item]}) + + +def test_estimate_effort_example_is_a_strict_integer(): + text = "estimated_effort_to_review_[1-5]: 3" + assert type(load_yaml(text)["estimated_effort_to_review_[1-5]"]) is int + Review.model_validate({ + "key_issues_to_review": [], + "estimated_effort_to_review_[1-5]": load_yaml(text)["estimated_effort_to_review_[1-5]"], + }) + + +def test_ranking_example_is_a_valid_numeric_payload(): + text = 'which_response_was_better: 1\nwhy: "It is clearer."\nscore_response1: 9\nscore_response2: 7' + parsed = load_yaml(text) + assert isinstance(parsed["which_response_was_better"], int) + PRRankResponses.model_validate(parsed) diff --git a/tests/unittest/test_per_directory_settings.py b/tests/unittest/test_per_directory_settings.py new file mode 100644 index 0000000000..5db49efdd0 --- /dev/null +++ b/tests/unittest/test_per_directory_settings.py @@ -0,0 +1,1236 @@ +import copy +from types import SimpleNamespace +from unittest.mock import MagicMock, call + +import pytest +from github import GithubException +from starlette_context import context, request_cycle_context + +from pr_agent.config_loader import get_settings, global_settings +from pr_agent.git_providers import utils as git_utils +from pr_agent.git_providers.github_provider import GithubProvider +from pr_agent.git_providers.gitlab_provider import GitLabProvider + +ROOT_TOML = b""" +[pr_reviewer] +num_max_findings = 10 + +[config] +model = "root-model" +temperature = 0.1 +""" + +SERVICES_TOML = b""" +[pr_reviewer] +num_max_findings = 5 + +[ignore] +glob = ["*.gen.js"] +""" + +SERVICES_AUTH_TOML = b""" +[pr_reviewer] +num_max_findings = 3 + +[config] +temperature = 0.5 +model = "auth-model" +""" + +SERVICES_BILLING_TOML = b""" +[pr_reviewer] +num_max_findings = 7 +""" + + +class FakePerDirProvider: + def __init__(self, root_settings=None, tree_paths=(), contents=None, files=None, + resolved_ref="main", full_files=None): + self.root_settings = root_settings + self.tree_paths = tuple(tree_paths) + self.contents = dict(contents) if contents else {} + self.files = list(files) if files else [] + self.full_files = list(full_files) if full_files is not None else None + self.resolved_ref = resolved_ref + self.tree_calls = 0 + self.tree_refs = [] + self.contents_calls = [] + self.get_files_calls = 0 + self.pr_file_paths_calls = 0 + + def get_repo_settings(self): + return self.root_settings + + def get_files(self): + self.get_files_calls += 1 + return self.files + + def get_pr_file_paths(self): + self.pr_file_paths_calls += 1 + return self.full_files if self.full_files is not None else self.files + + def get_repo_settings_tree(self, ref): + self.tree_calls += 1 + self.tree_refs.append(ref) + return list(self.tree_paths), self.resolved_ref + + def get_repo_settings_contents(self, paths, ref): + self.contents_calls.append((list(paths), ref)) + return {path: self.contents[path] for path in paths if path in self.contents} + + def is_supported(self, capability): + return False + + def publish_comment(self, body): + pass + + def publish_persistent_comment(self, *args, **kwargs): + pass + + +@pytest.fixture +def fresh_global_settings(): + """Restore module-level global_settings after each test in case anything mutated it.""" + snapshot = copy.deepcopy(global_settings.as_dict()) + yield + for section in set(global_settings.as_dict().keys()) - set(snapshot.keys()): + global_settings.unset(section) + for section, contents in snapshot.items(): + global_settings.unset(section) + global_settings.set(section, copy.deepcopy(contents), merge=False) + + +@pytest.fixture +def per_dir_settings(fresh_global_settings): + """Request-scoped settings clone with the per-directory feature enabled.""" + with request_cycle_context({}): + context["settings"] = copy.deepcopy(global_settings) + settings = get_settings() + settings.config.enable_per_directory_settings = True + yield + + +def _provider(tree_paths, contents, files, root_settings=b"", full_files=None): + return FakePerDirProvider( + root_settings=root_settings, + tree_paths=tree_paths, + contents=contents, + files=files, + full_files=full_files, + ) + + +class TestResolvePerDirectorySettings: + def test_literal_backslashes_do_not_activate_nested_settings(self, per_dir_settings): + provider = _provider( + tree_paths=["services/auth/.pr_agent.toml"], + contents={"services/auth/.pr_agent.toml": SERVICES_AUTH_TOML}, + files=[r"services\auth\file.py"], + ) + + assert git_utils._get_per_directory_settings(provider) == [] + assert provider.contents_calls == [] + + def test_walks_up_from_changed_file_to_root(self, per_dir_settings): + provider = _provider( + tree_paths=[".pr_agent.toml", "services/.pr_agent.toml", "services/auth/.pr_agent.toml"], + contents={ + "services/.pr_agent.toml": SERVICES_TOML, + "services/auth/.pr_agent.toml": SERVICES_AUTH_TOML, + }, + files=["services/auth/api.py", "other/README.md"], + ) + + resolved = git_utils._get_per_directory_settings(provider) + + # The root .pr_agent.toml is already applied through get_repo_settings() and + # must not be re-applied here; ancestor configs are returned shallowest-first + # so a nearer file overrides a farther one. + assert [path for path, _ in resolved] == [ + "services/.pr_agent.toml", + "services/auth/.pr_agent.toml", + ] + + def test_sibling_configs_both_applied(self, per_dir_settings): + provider = _provider( + tree_paths=["services/auth/.pr_agent.toml", "services/billing/.pr_agent.toml"], + contents={ + "services/auth/.pr_agent.toml": SERVICES_AUTH_TOML, + "services/billing/.pr_agent.toml": SERVICES_BILLING_TOML, + }, + files=["services/auth/api.py", "services/billing/x.py"], + ) + + resolved = git_utils._get_per_directory_settings(provider) + + assert sorted(path for path, _ in resolved) == [ + "services/auth/.pr_agent.toml", + "services/billing/.pr_agent.toml", + ] + + def test_sibling_overlap_is_detected_and_warned(self): + from loguru import logger as loguru_logger + + ordered = ["services/auth", "services/billing"] + contents = { + "services/auth/.pr_agent.toml": SERVICES_AUTH_TOML, + "services/billing/.pr_agent.toml": SERVICES_BILLING_TOML, + } + + captured_lines = [] + sink_id = loguru_logger.add( + lambda msg: captured_lines.append(str(msg)), + level="WARNING", + ) + try: + conflicts = git_utils._warn_on_sibling_key_conflicts(ordered, contents) + finally: + loguru_logger.remove(sink_id) + + assert conflicts == [("pr_reviewer", "num_max_findings", "services/billing")] + assert any("pr_reviewer.num_max_findings" in line for line in captured_lines) + + def test_sibling_conflict_ignores_disjoint_keys(self): + ordered = ["services/auth", "services/billing"] + contents = { + "services/auth/.pr_agent.toml": SERVICES_AUTH_TOML, + "services/billing/.pr_agent.toml": b"[pr_description]\nuse_description_markers = true\n", + } + + conflicts = git_utils._warn_on_sibling_key_conflicts(ordered, contents) + + assert conflicts == [] + + def test_sibling_key_shared_by_two_of_three_is_warned(self): + from loguru import logger as loguru_logger + + ordered = ["services/auth", "services/billing", "services/orders"] + contents = { + "services/auth/.pr_agent.toml": SERVICES_AUTH_TOML, + "services/billing/.pr_agent.toml": SERVICES_BILLING_TOML, + "services/orders/.pr_agent.toml": b"[pr_description]\nuse_description_markers = true\n", + } + + captured_lines = [] + sink_id = loguru_logger.add( + lambda msg: captured_lines.append(str(msg)), + level="WARNING", + ) + try: + conflicts = git_utils._warn_on_sibling_key_conflicts(ordered, contents) + finally: + loguru_logger.remove(sink_id) + + # pr_reviewer.num_max_findings is set by auth + billing only; orders only sets config.temperature. + assert conflicts == [("pr_reviewer", "num_max_findings", "services/billing")] + assert any("pr_reviewer.num_max_findings" in line for line in captured_lines) + + def test_sibling_winner_is_last_owner_that_sets_key(self): + ordered = ["services/auth", "services/billing", "services/orders"] + contents = { + "services/auth/.pr_agent.toml": SERVICES_AUTH_TOML, + "services/billing/.pr_agent.toml": SERVICES_BILLING_TOML, + # orders is the last sibling in path order but its content is invalid TOML, + # so it can never be the winner for a key it does not set. + "services/orders/.pr_agent.toml": b"[pr_reviewer\nbroken = true\n", + } + + conflicts = git_utils._warn_on_sibling_key_conflicts(ordered, contents) + + assert conflicts == [("pr_reviewer", "num_max_findings", "services/billing")] + + def test_no_config_crossed_returns_empty(self, per_dir_settings): + provider = _provider( + tree_paths=["services/auth/.pr_agent.toml"], + contents={"services/auth/.pr_agent.toml": SERVICES_AUTH_TOML}, + files=["docs/readme.md"], + ) + + assert git_utils._get_per_directory_settings(provider) == [] + + def test_cap_keeps_shallowest_configs(self, per_dir_settings): + get_settings().config.per_directory_settings_max_files = 2 + provider = _provider( + tree_paths=[ + "svc1/.pr_agent.toml", + "svc1/deep/.pr_agent.toml", + "svc2/.pr_agent.toml", + ], + contents={ + "svc1/.pr_agent.toml": SERVICES_TOML, + "svc1/deep/.pr_agent.toml": SERVICES_AUTH_TOML, + "svc2/.pr_agent.toml": SERVICES_BILLING_TOML, + }, + files=["svc1/deep/api.py", "svc2/x.py"], + ) + + resolved = git_utils._get_per_directory_settings(provider) + + # Shallowest first: the deeper svc1/deep config is dropped once the cap is full. + assert [path for path, _ in resolved] == [ + "svc1/.pr_agent.toml", + "svc2/.pr_agent.toml", + ] + + def test_cap_cutting_an_equal_depth_group_keeps_laterpath_winners(self, per_dir_settings): + get_settings().config.per_directory_settings_max_files = 2 + provider = _provider( + tree_paths=[ + "svc1/.pr_agent.toml", + "svc2/.pr_agent.toml", + "svc3/.pr_agent.toml", + ], + contents={ + "svc1/.pr_agent.toml": SERVICES_TOML, + "svc2/.pr_agent.toml": SERVICES_BILLING_TOML, + "svc3/.pr_agent.toml": SERVICES_AUTH_TOML, + }, + files=["svc1/x.py", "svc2/y.py", "svc3/z.py"], + ) + + resolved = git_utils._get_per_directory_settings(provider) + + # The cap cuts through an equal-depth group: the lexicographically-later (winning) + # siblings are retained and svc3 still participates as the documented winner. + assert [path for path, _ in resolved] == [ + "svc2/.pr_agent.toml", + "svc3/.pr_agent.toml", + ] + + def test_per_directory_oversized_config_is_dropped(self, per_dir_settings, monkeypatch): + monkeypatch.setattr(git_utils, "MAX_TOML_SIZE_IN_BYTES", 100) + provider = _provider( + tree_paths=["svc1/.pr_agent.toml", "svc2/.pr_agent.toml"], + contents={ + "svc1/.pr_agent.toml": b"[config]\n" + b"#" * 200, + "svc2/.pr_agent.toml": SERVICES_TOML, + }, + files=["svc1/x.py", "svc2/y.py"], + ) + + resolved = git_utils._get_per_directory_settings(provider) + + # An over-limit nested file is dropped before retention or parsing; the valid one stays. + assert [path for path, _ in resolved] == ["svc2/.pr_agent.toml"] + + def test_per_directory_aggregate_budget_stops_retention(self, per_dir_settings, monkeypatch): + monkeypatch.setattr( + git_utils, "_PER_DIRECTORY_SETTINGS_AGGREGATE_BYTES", len(SERVICES_TOML) + 10 + ) + provider = _provider( + tree_paths=[ + "svc1/.pr_agent.toml", + "svc2/.pr_agent.toml", + "svc3/.pr_agent.toml", + ], + contents={ + "svc1/.pr_agent.toml": SERVICES_TOML, + "svc2/.pr_agent.toml": SERVICES_BILLING_TOML, + "svc3/.pr_agent.toml": SERVICES_AUTH_TOML, + }, + files=["svc1/x.py", "svc2/y.py", "svc3/z.py"], + ) + + reserved = git_utils._get_per_directory_settings(provider) + + # Once the aggregate budget is exhausted, further files are not retained. + assert [path for path, _ in reserved] == ["svc1/.pr_agent.toml"] + + def test_sibling_conflict_parse_skips_oversized_files(self, per_dir_settings, monkeypatch): + monkeypatch.setattr(git_utils, "MAX_TOML_SIZE_IN_BYTES", 100) + oversized = b"[pr_reviewer]\nnum_max_findings = 10\n" + b"#" * 300 + contents = { + "svc1/.pr_agent.toml": oversized, + "svc2/.pr_agent.toml": b"[pr_reviewer]\nnum_max_findings = 5\n", + } + + conflicts = git_utils._warn_on_sibling_key_conflicts(["svc1", "svc2"], contents) + + # The oversized sibling is not parsed, so no conflict with svc2 can be reported. + assert conflicts == [] + + def test_disabled_returns_empty_without_any_provider_call(self, fresh_global_settings): + with request_cycle_context({}): + context["settings"] = copy.deepcopy(global_settings) + provider = _provider( + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": SERVICES_TOML}, + files=["services/api.py"], + ) + + assert git_utils._get_per_directory_settings(provider) == [] + assert provider.tree_calls == 0 + assert provider.get_files_calls == 0 + + def test_provider_without_support_is_inert(self, per_dir_settings): + class BareProvider: + def __init__(self): + self.get_files_calls = 0 + + def get_repo_settings(self): + return b"" + + def get_files(self): + self.get_files_calls += 1 + return ["services/api.py"] + + provider = BareProvider() + assert git_utils._get_per_directory_settings(provider) == [] + assert provider.get_files_calls == 0 + + def test_changed_file_path_normalization(self, per_dir_settings): + provider = _provider( + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": SERVICES_TOML}, + files=[ + "services/plain.py", + {"new_path": "services/dict_new.py"}, + {"filename": "services/dict_filename.py"}, + SimpleNamespace(filename="services/obj_filename.py"), + ], + ) + + resolved = git_utils._get_per_directory_settings(provider) + + assert [path for path, _ in resolved] == ["services/.pr_agent.toml"] + + def test_changed_file_paths_include_rename_metadata(self, per_dir_settings): + provider = _provider( + tree_paths=["legacy/.pr_agent.toml", "services/.pr_agent.toml"], + contents={ + "legacy/.pr_agent.toml": SERVICES_TOML, + "services/.pr_agent.toml": SERVICES_TOML, + }, + files=[ + {"new_path": "services/api.py", "old_path": "legacy/api.py"}, + SimpleNamespace(new_path="services/web.py", previous_filename="legacy/web.py"), + {"new_path": "services/fresh.py", "old_path": None}, + ], + ) + + resolved = git_utils._get_per_directory_settings(provider) + + # Both the source and the destination side of each move contribute paths, + # so their ancestor configs apply; a bare new file contributes only its own. + assert {path for path, _ in resolved} == { + "legacy/.pr_agent.toml", + "services/.pr_agent.toml", + } + + def test_incremental_state_does_not_shrink_discovery(self, per_dir_settings): + provider = _provider( + tree_paths=["services/.pr_agent.toml", "services/auth/.pr_agent.toml"], + contents={ + "services/.pr_agent.toml": SERVICES_TOML, + "services/auth/.pr_agent.toml": SERVICES_AUTH_TOML, + }, + files=["services/api.py"], + full_files=["services/api.py", "services/auth/other.py"], + ) + + resolved = git_utils._get_per_directory_settings(provider) + + # get_files() narrows to the unreviewed subset once an incremental review + # is active; discovery must use the complete PR file set instead, so a later + # command in the same request does not silently drop per-directory configs. + assert {path for path, _ in resolved} == { + "services/.pr_agent.toml", + "services/auth/.pr_agent.toml", + } + assert provider.get_files_calls == 0 + assert provider.pr_file_paths_calls == 1 + + def test_config_branch_passed_to_tree(self, per_dir_settings): + get_settings().set("CONFIG.CONFIG_BRANCH", "cfg-branch") + provider = _provider( + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": SERVICES_TOML}, + files=["services/api.py"], + ) + + git_utils._get_per_directory_settings(provider) + + assert provider.tree_refs == ["cfg-branch"] + + def test_get_files_failure_degrades_to_empty(self, per_dir_settings): + class ExplodingProvider(FakePerDirProvider): + def get_pr_file_paths(self): + raise RuntimeError("boom") + + provider = ExplodingProvider( + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": SERVICES_TOML}, + ) + + assert git_utils._get_per_directory_settings(provider) == [] + + def test_tree_failure_degrades_to_empty(self, per_dir_settings): + class ExplodingProvider(FakePerDirProvider): + def get_repo_settings_tree(self, ref): + raise RuntimeError("boom") + + provider = ExplodingProvider(files=["services/api.py"]) + + assert git_utils._get_per_directory_settings(provider) == [] + + +class TestApplyPerDirectorySettings: + def test_merge_root_then_directory_nearest_wins(self, per_dir_settings, monkeypatch): + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml", "services/auth/.pr_agent.toml"], + contents={ + "services/.pr_agent.toml": SERVICES_TOML, + "services/auth/.pr_agent.toml": SERVICES_AUTH_TOML, + }, + files=["services/auth/api.py", "services/billing/x.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + # Nearest (services/auth) wins over services, which wins over root. + assert get_settings().pr_reviewer.num_max_findings == 3 + assert get_settings().config.temperature == 0.5 + assert get_settings().config.model == "auth-model" + + def test_different_key_casing_overrides_root_value(self, per_dir_settings, monkeypatch): + # Dynaconf resolves keys case-insensitively, so a per-directory file that spells + # the same key with different casing must replace the root value, not join it as a + # duplicate that leaves the old value winning (regression for "nearest wins"). + nested = b""" +[pr_reviewer] +Num_Max_Findings = 4 +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=b"[pr_reviewer]\nnum_max_findings = 10\n", + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": nested}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + assert get_settings().pr_reviewer.num_max_findings == 4 + + def test_list_values_replace_not_concatenate(self, per_dir_settings, monkeypatch): + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": SERVICES_TOML}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + # The default [ignore] glob ('vendor/**') must be replaced by the per-directory + # file wholesale, not concatenated with it. + assert get_settings().ignore.glob == ["*.gen.js"] + + def test_whitelist_blocks_secrets_and_critical_sections(self, per_dir_settings, monkeypatch): + evil = b""" +[openai] +api_base = "https://evil.example.com" + +[config] +model = "allowed-model" +git_provider = "gitlab" + +[push_outputs] +webhook_url = "https://evil.example.com/hook" + +[pr_reviewer] +num_max_findings = 2 +publish_error_details = true +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": evil}, + files=["services/api.py"], + ), + ) + # Baseline: ensure inherited settings from earlier tests don't mask a leak. + get_settings().config.git_provider = "github" + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + # [openai], [push_outputs] and config.git_provider are not per-directory overridable. + assert get_settings().get("openai.api_base", None) != "https://evil.example.com" + assert get_settings().config.git_provider != "gitlab" + assert get_settings().get("push_outputs.webhook_url", None) != "https://evil.example.com/hook" + # Allowed keys still land, and the repo-host-only pr_reviewer key stays dropped. + assert get_settings().config.model == "allowed-model" + assert get_settings().pr_reviewer.num_max_findings == 2 + assert get_settings().pr_reviewer.publish_error_details is False + + def test_whitelist_rejects_unknown_sections_entirely(self, per_dir_settings, monkeypatch): + config = b""" +[openai] +api_base = "https://evil.example.com" + +[pr_reviewer] +num_max_findings = 4 +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=b"", + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + assert get_settings().get("openai.api_base", None) != "https://evil.example.com" + assert get_settings().pr_reviewer.num_max_findings == 4 + + def test_write_or_url_trigger_keys_are_dropped_in_per_directory(self, per_dir_settings, monkeypatch): + config = b""" +[pr_update_changelog] +push_changelog_changes = true +add_pr_link = true +extra_instructions = "keep-me" + +[pr_help_docs] +repo_url = "https://evil.example.com/steal" +docs_path = "custom-docs" +exclude_root_readme = true +supported_doc_exts = [".md"] +enable_help_text = true +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + assert get_settings().pr_update_changelog.push_changelog_changes is False + assert get_settings().get("pr_help_docs.repo_url", "") == "" + # Collection scope is root-/host-controlled: a nested file cannot point /help_docs + # at arbitrary repository paths or file extensions to read into the model prompt. + assert get_settings().pr_help_docs.docs_path == "docs" + assert get_settings().pr_help_docs.supported_doc_exts == [".md", ".mdx", ".rst"] + assert get_settings().pr_update_changelog.add_pr_link is True + assert get_settings().pr_update_changelog.extra_instructions == "keep-me" + assert get_settings().pr_help_docs.exclude_root_readme is True + assert get_settings().pr_help_docs.enable_help_text is True + + def test_description_questions_and_similar_issue_host_only_keys_are_dropped( + self, per_dir_settings, monkeypatch + ): + config = b""" +[pr_reviewer] +num_max_findings = 4 +enable_large_pr_chunking = true +max_number_of_calls = 9999 +inline_key_issues = true +enable_review_labels_security = false +enable_review_labels_effort = false +require_security_review = false +require_estimate_effort_to_review = false +require_ticket_analysis_review = false + +[pr_description] +publish_labels = true +use_ai_title = true +generate_ai_title = true +publish_description_as_comment = true +publish_description_as_comment_persistent = false +enable_large_pr_handling = true +max_ai_calls = 400 +async_ai_calls = false + +[pr_questions] +resolve_threads = true +static_questions = ["default"] +use_conversation_history = true + +[pr_code_suggestions] +commitable_code_suggestions = true +num_code_suggestions_per_chunk = 2 +max_number_of_calls = 99 +parallel_calls = true +demand_code_suggestions_self_review = true +approve_pr_on_self_review = true + +[pr_similar_issue] +force_update_dataset = true +skip_comments = false +max_issues_to_scan = 999999 +vectordb = "pinecone" +use_original_title = false +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML + b""" +[pr_questions] +use_conversation_history = false +[pr_similar_issue] +skip_comments = true +""", + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + # Label mutation, thread resolution and full issue-index refresh switches are + # host-/root-controlled and must stay at their trusted defaults. + assert get_settings().pr_description.publish_labels is False + assert get_settings().pr_questions.resolve_threads is False + assert get_settings().pr_similar_issue.force_update_dataset is False + assert get_settings().pr_similar_issue.skip_comments is True + assert get_settings().get("pr_similar_issue.max_issues_to_scan", 0) != 999999 + assert get_settings().pr_similar_issue.vectordb != "pinecone" + # Pull-request metadata and reviewer label/inline controls stay trusted: + # a nested file cannot rewrite PR titles/bodies, publish reviewer effort or + # security labels, or open inline key-issue comments through the bot identity. + assert get_settings().pr_description.generate_ai_title is False + assert get_settings().pr_description.publish_description_as_comment is False + assert get_settings().pr_description.publish_description_as_comment_persistent is True + assert get_settings().pr_reviewer.inline_key_issues is False + assert get_settings().pr_reviewer.enable_review_labels_security is True + assert get_settings().pr_reviewer.enable_review_labels_effort is True + assert get_settings().pr_reviewer.require_security_review is True + assert get_settings().pr_reviewer.require_estimate_effort_to_review is True + assert get_settings().pr_reviewer.require_ticket_analysis_review is True + # Budget/call-count controls stay at the host-trusted values: nested files must + # not multiply AI calls on their own. + assert get_settings().pr_reviewer.enable_large_pr_chunking is False + assert get_settings().pr_reviewer.max_number_of_calls != 9999 + assert get_settings().pr_description.enable_large_pr_handling != 999999 + assert get_settings().pr_description.max_ai_calls != 400 + assert get_settings().pr_description.async_ai_calls is True + assert get_settings().pr_code_suggestions.max_number_of_calls != 99 + assert get_settings().pr_code_suggestions.parallel_calls is True + # The self-review approval workflow stays root-controlled: a nested file cannot + # demand a self-review checklist and then auto-approve on the author's tick. + assert get_settings().pr_code_suggestions.demand_code_suggestions_self_review is False + assert get_settings().pr_code_suggestions.approve_pr_on_self_review is False + # Thread-history collection for /ask is root-controlled: a nested file cannot + # re-enable sending private review-thread discussion bodies to the model. + assert get_settings().pr_questions.use_conversation_history is False + # Ordinary keys in the same sections still apply. + assert get_settings().pr_reviewer.num_max_findings == 4 + assert get_settings().pr_description.use_ai_title is True + assert get_settings().pr_questions.static_questions == ["default"] + assert get_settings().pr_code_suggestions.commitable_code_suggestions is False + assert get_settings().pr_code_suggestions.num_code_suggestions_per_chunk == 2 + assert get_settings().pr_similar_issue.use_original_title is False + + def test_per_directory_ignore_allows_only_glob(self, per_dir_settings, monkeypatch): + config = b""" +[ignore] +glob = ["gen/**"] +regex = ["(a+)+$"] +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + # Bounded glob patterns land; arbitrary regexes (which filter_ignored() compiles + # and matches on every review) stay at the root default so a nested file cannot + # commit a catastrophic-backtracking pattern. + assert get_settings().ignore.glob == ["gen/**"] + assert get_settings().ignore.regex == [] + + def test_per_directory_config_drops_unbounded_repo_context_knobs(self, per_dir_settings, monkeypatch): + config = b""" +[config] +model = "nested-model" +temperature = 0.5 +repo_context_files = ["huge.bin", "secrets.env", "vendor/data.bin"] +repo_context_max_lines = 9999999 +model_token_count_estimate_factor = -0.999999 +per_directory_settings_max_tree_pages = 999999 +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + # Context-fetch knobs are root-/host-controlled; a nested file cannot make tools + # fetch arbitrary repo files in full when they build context for its directory. + assert get_settings().get("config.repo_context_files", []) != [ + "huge.bin", + "secrets.env", + "vendor/data.bin", + ] + assert get_settings().get("config.repo_context_max_lines", 0) != 9999999 + assert get_settings().config.model_token_count_estimate_factor == 0.3 + assert get_settings().config.per_directory_settings_max_tree_pages == 10 + # Model-routing and output knobs still apply. + assert get_settings().config.model == "nested-model" + assert get_settings().config.temperature == 0.5 + + def test_per_directory_config_drops_non_string_model_values(self, per_dir_settings, monkeypatch): + trusted_model_weak = get_settings().get("config.model_weak", None) + config = b""" +[config] +model = 1 +model_weak = false +model_reasoning = "reasoning-model" +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + assert get_settings().config.model == "root-model" + assert get_settings().get("config.model_weak", None) == trusted_model_weak + assert get_settings().config.model_reasoning == "reasoning-model" + + @pytest.mark.parametrize("invalid_value", [(1, "1"), (False, "false"), (["fr-fr"], '["fr-fr"]'), + ({"locale": "fr-fr"}, '{locale = "fr-fr"}')]) + def test_per_directory_config_drops_non_string_response_language(self, per_dir_settings, monkeypatch, + invalid_value): + _, toml_value = invalid_value + config = f"[config]\nresponse_language = {toml_value}\n".encode() + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + assert get_settings().get("config.response_language", "en-us") == "en-US" + + @pytest.mark.parametrize("toml_value", [ + '"hot"', "true", "[0.1]", '{value = "0.1"}', "nan", "-0.1", "2.1", + ]) + def test_per_directory_config_drops_invalid_temperature(self, per_dir_settings, monkeypatch, toml_value): + config = f"[config]\ntemperature = {toml_value}\n".encode() + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + assert get_settings().config.temperature == 0.1 + + @pytest.mark.parametrize("toml_value", ['"six"', "true", "-1", "1001", "1.5"]) + def test_per_directory_config_drops_invalid_description_threshold(self, per_dir_settings, monkeypatch, + toml_value): + config = f"[pr_description]\ncollapsible_file_list_threshold = {toml_value}\n".encode() + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + assert get_settings().get("pr_description.collapsible_file_list_threshold") == 6 + + def test_per_directory_config_drops_fallback_models(self, per_dir_settings, monkeypatch): + config = b""" +[config] +model = "nested-model" +temperature = 0.5 +fallback_models = ["m-one", "m-two", "m-three", "m-four", "m-five"] +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + # Fallback routing stays root-/host-controlled: the retry helper turns each + # entry into one more routed attempt per failing model, so a nested file must + # not be able to multiply AI calls with an arbitrarily long list. + assert list(get_settings().config.fallback_models) == ["gpt-5.6-terra"] + # Other model-routing and output knobs still apply. + assert get_settings().config.model == "nested-model" + assert get_settings().config.temperature == 0.5 + + def test_per_directory_config_drops_token_budget_keys(self, per_dir_settings, monkeypatch): + config = b""" +[config] +model = "nested-model" +temperature = 0.5 +max_model_tokens = 999999 +custom_model_max_tokens = 999999 +max_output_tokens = 999999 +""" + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": config}, + files=["services/api.py"], + ), + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + # Token budgets are root-/host-controlled because they directly size request + # and completion limits; a nested file cannot inflate the cost of every call. + assert get_settings().config.max_model_tokens != 999999 + assert get_settings().config.custom_model_max_tokens != 999999 + assert get_settings().config.max_output_tokens != 999999 + # Model-routing and output knobs still apply. + assert get_settings().config.model == "nested-model" + assert get_settings().config.temperature == 0.5 + + def test_malformed_per_directory_config_reports_error(self, per_dir_settings, monkeypatch): + malformed = b"[pr_reviewer\nnum_max_findings = 2\n" + provider = _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": malformed}, + files=["services/api.py"], + ) + comments = [] + provider.publish_persistent_comment = lambda *args, **kwargs: comments.append((args, kwargs)) + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: provider, + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + assert len(comments) == 1 + assert "services/.pr_agent.toml" in comments[0][0][0] + # The root config still applied before the malformed per-directory file. + assert get_settings().pr_reviewer.num_max_findings == 10 + + def test_per_directory_inert_when_feature_disabled(self, fresh_global_settings, monkeypatch): + with request_cycle_context({}): + context["settings"] = copy.deepcopy(global_settings) + provider = _provider( + root_settings=ROOT_TOML, + tree_paths=["services/.pr_agent.toml"], + contents={"services/.pr_agent.toml": SERVICES_AUTH_TOML}, + files=["services/auth/api.py"], + ) + monkeypatch.setattr( + "pr_agent.git_providers.utils.get_git_provider_with_context", + lambda url: provider, + ) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + + assert provider.tree_calls == 0 + assert provider.get_files_calls == 0 + assert get_settings().pr_reviewer.num_max_findings == 10 + + +def _github_provider(repo_obj): + provider = GithubProvider.__new__(GithubProvider) + provider.repo_obj = repo_obj + provider._resolved_config_branch = None + return provider + + +class TestGithubProviderPerDirectory: + def test_get_repo_settings_tree_filters_pr_agent_toml_blobs(self): + repo_obj = MagicMock() + repo_obj.default_branch = "main" + repo_obj.get_git_tree.return_value = SimpleNamespace(tree=[ + SimpleNamespace(path=".pr_agent.toml", type="blob"), + SimpleNamespace(path="services/auth/.pr_agent.toml", type="blob"), + SimpleNamespace(path="services/auth/code.py", type="blob"), + SimpleNamespace(path="services/auth", type="tree"), + ]) + provider = _github_provider(repo_obj) + + paths, resolved_ref = provider.get_repo_settings_tree() + + assert resolved_ref == "main" + assert paths == [".pr_agent.toml", "services/auth/.pr_agent.toml"] + repo_obj.get_git_tree.assert_called_once_with("main", recursive=True) + + def test_get_repo_settings_tree_uses_resolved_config_branch(self): + repo_obj = MagicMock() + repo_obj.get_git_tree.return_value = SimpleNamespace(tree=[]) + provider = _github_provider(repo_obj) + provider._resolved_config_branch = "cfg-branch" + + _, resolved_ref = provider.get_repo_settings_tree("") + + assert resolved_ref == "cfg-branch" + repo_obj.get_git_tree.assert_called_once_with("cfg-branch", recursive=True) + + def test_resolved_config_branch_beats_explicit_ref(self): + # get_repo_settings() stores the branch it actually read the root config + # from (already fallback-resolved), so it must win over a CONFIG_BRANCH hint: + # when that branch exists without a root .pr_agent.toml the tree must follow + # the root config onto the default branch instead of reading a stale branch. + repo_obj = MagicMock() + repo_obj.get_git_tree.return_value = SimpleNamespace(tree=[], truncated=False) + provider = _github_provider(repo_obj) + provider._resolved_config_branch = "resolved-default" + + _, resolved_ref = provider.get_repo_settings_tree("stale-config-branch") + + assert resolved_ref == "resolved-default" + repo_obj.get_git_tree.assert_called_once_with("resolved-default", recursive=True) + + def test_truncated_tree_skips_per_directory_settings(self): + from loguru import logger as loguru_logger + + repo_obj = MagicMock() + repo_obj.get_git_tree.return_value = SimpleNamespace( + tree=[SimpleNamespace(path="svc/.pr_agent.toml", type="blob")], + truncated=True, + ) + provider = _github_provider(repo_obj) + + captured_lines = [] + sink_id = loguru_logger.add( + lambda msg: captured_lines.append(str(msg)), + level="WARNING", + ) + try: + paths, resolved_ref = provider.get_repo_settings_tree("big-branch") + finally: + loguru_logger.remove(sink_id) + + assert paths == [] + assert resolved_ref == "big-branch" + assert any("truncated" in line for line in captured_lines) + + def test_get_repo_settings_tree_falls_back_to_default_on_404(self): + repo_obj = MagicMock() + repo_obj.default_branch = "main" + repo_obj.get_git_tree.side_effect = [ + GithubException(404, {"message": "Not Found"}, None), + SimpleNamespace(tree=[ + SimpleNamespace(path="svc/.pr_agent.toml", type="blob"), + ]), + ] + provider = _github_provider(repo_obj) + + paths, resolved_ref = provider.get_repo_settings_tree("missing-branch") + + assert resolved_ref == "main" + assert paths == ["svc/.pr_agent.toml"] + assert repo_obj.get_git_tree.call_args_list == [call("missing-branch", recursive=True), call("main", recursive=True)] + + def test_get_repo_settings_tree_surfaces_unexpected_errors(self): + repo_obj = MagicMock() + repo_obj.get_git_tree.side_effect = GithubException(403, {"message": "Forbidden"}, None) + provider = _github_provider(repo_obj) + + with pytest.raises(GithubException): + provider.get_repo_settings_tree("") + + def test_get_repo_settings_contents(self): + repo_obj = MagicMock() + repo_obj.get_contents.return_value = SimpleNamespace(decoded_content=b"[pr_reviewer]\nnum_max_findings = 5\n") + provider = _github_provider(repo_obj) + + result = provider.get_repo_settings_contents(["services/auth/.pr_agent.toml"], "main") + + assert result == {"services/auth/.pr_agent.toml": b"[pr_reviewer]\nnum_max_findings = 5\n"} + repo_obj.get_contents.assert_called_once_with("services/auth/.pr_agent.toml", ref="main") + + def test_get_repo_settings_contents_skips_missing_file(self): + repo_obj = MagicMock() + repo_obj.get_contents.side_effect = GithubException(404, {"message": "Not Found"}, None) + provider = _github_provider(repo_obj) + + result = provider.get_repo_settings_contents(["services/.pr_agent.toml"], "main") + + assert result == {} + + def test_get_pr_file_paths_returns_full_set_with_rename_metadata(self): + full_files = [ + SimpleNamespace(filename="services/api.py", previous_filename="legacy/api.py"), + SimpleNamespace(filename="services/fresh.py", previous_filename=None), + ] + provider = _github_provider(MagicMock()) + provider.git_files = full_files + provider.incremental = SimpleNamespace(is_incremental=True) + provider.unreviewed_files_map = {"services/api.py": "bogus-subset"} + + result = provider.get_pr_file_paths() + + # Even with an incremental review active (unreviewed_files_map populated), the + # complete PR file set with rename metadata is returned, never the reviewed subset. + assert result == full_files + assert [entry.previous_filename for entry in result] == ["legacy/api.py", None] + + +def _gitlab_provider(gl, id_project="1"): + provider = GitLabProvider.__new__(GitLabProvider) + provider.gl = gl + provider.id_project = id_project + return provider + + +class TestGitLabProviderPerDirectory: + def test_get_repo_settings_tree_filters_blobs_on_default_branch(self): + project = MagicMock() + project.default_branch = "main" + project.repository_tree.return_value = [ + {"path": ".pr_agent.toml", "type": "blob"}, + {"path": "svc/.pr_agent.toml", "type": "blob"}, + {"path": "svc/code.py", "type": "blob"}, + {"path": "svc", "type": "tree"}, + ] + gl = MagicMock() + gl.projects.get.return_value = project + provider = _gitlab_provider(gl) + + paths, resolved_ref = provider.get_repo_settings_tree("ignored-ref") + + assert resolved_ref == "main" + assert paths == [".pr_agent.toml", "svc/.pr_agent.toml"] + project.repository_tree.assert_called_once_with(ref="main", recursive=True, page=1, per_page=100) + + @pytest.mark.parametrize("complete", [True, False]) + def test_tree_discovery_is_bounded_and_requires_complete_results(self, per_dir_settings, complete): + get_settings().config.per_directory_settings_max_tree_pages = 2 + project = MagicMock() + project.default_branch = "main" + full_page = [{"path": "svc/.pr_agent.toml", "type": "blob"}] * 100 + project.repository_tree.side_effect = [full_page, [] if complete else full_page] + gl = MagicMock() + gl.projects.get.return_value = project + + paths, ref = _gitlab_provider(gl).get_repo_settings_tree() + + assert ref == "main" + assert paths == (["svc/.pr_agent.toml"] * 100 if complete else []) + assert project.repository_tree.call_args_list == [ + call(ref="main", recursive=True, page=1, per_page=100), + call(ref="main", recursive=True, page=2, per_page=100), + ] + + def test_get_repo_settings_contents(self): + project = MagicMock() + file_obj = MagicMock() + file_obj.decode.return_value = b"[pr_reviewer]\nnum_max_findings = 5\n" + project.files.get.return_value = file_obj + gl = MagicMock() + gl.projects.get.return_value = project + provider = _gitlab_provider(gl) + + result = provider.get_repo_settings_contents(["svc/.pr_agent.toml"], "main") + + assert result == {"svc/.pr_agent.toml": b"[pr_reviewer]\nnum_max_findings = 5\n"} + project.files.get.assert_called_once_with(file_path="svc/.pr_agent.toml", ref="main") + + def test_get_repo_settings_contents_skips_missing_file(self): + from gitlab.exceptions import GitlabGetError + + project = MagicMock() + project.files.get.side_effect = GitlabGetError(response_code=404) + gl = MagicMock() + gl.projects.get.return_value = project + provider = _gitlab_provider(gl) + + result = provider.get_repo_settings_contents(["svc/.pr_agent.toml"], "main") + + assert result == {} + + def test_get_pr_file_paths_returns_full_changes_with_rename_metadata(self): + provider = _gitlab_provider(MagicMock()) + changes = [ + {"new_path": "services/api.py", "old_path": "legacy/api.py"}, + {"new_path": "services/fresh.py", "old_path": "services/fresh.py"}, + ] + provider._get_merge_request_changes = MagicMock(return_value={"changes": changes}) + provider._expand_submodule_changes = lambda ch: ch + + result = provider.get_pr_file_paths() + + # Both sides of a rename survive, independent of incremental review state. + assert result == changes + + +@pytest.mark.parametrize("next_root", [b"", b"[pr_reviewer]\nnum_max_findings = 12\n"]) +def test_directory_overrides_do_not_leak_between_commands(fresh_global_settings, monkeypatch, next_root): + settings = get_settings() + settings.config.enable_per_directory_settings = True + settings.config.use_repo_settings_file = True + settings.pr_reviewer.num_max_findings = 10 + providers = iter([ + _provider( + tree_paths=["services/.pr_agent.toml", "services/auth/.pr_agent.toml"], + contents={ + "services/.pr_agent.toml": SERVICES_TOML, + "services/auth/.pr_agent.toml": b"[pr_reviewer]\nnum_max_findings = 3\nnew_test_key = 42\n", + }, + files=["services/auth/api.py"], + ), + _provider(tree_paths=[], contents={}, files=["README.md"], root_settings=next_root), + ]) + monkeypatch.setattr(git_utils, "get_git_provider_with_context", lambda url: next(providers)) + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/1") + assert settings.pr_reviewer.num_max_findings == 3 + assert settings.pr_reviewer.new_test_key == 42 + settings.pr_reviewer.extra_instructions = "later trusted change" + + git_utils.apply_repo_settings("https://github.com/org/repo/pull/2") + assert settings.pr_reviewer.num_max_findings == (12 if next_root else 10) + assert "new_test_key" not in settings.pr_reviewer + assert settings.pr_reviewer.extra_instructions == "later trusted change" diff --git a/tests/unittest/test_plain_diff_provider.py b/tests/unittest/test_plain_diff_provider.py index 2f2f923be9..9c9d9c7a23 100644 --- a/tests/unittest/test_plain_diff_provider.py +++ b/tests/unittest/test_plain_diff_provider.py @@ -4,6 +4,7 @@ from pr_agent.config_loader import get_settings from pr_agent.git_providers import _GIT_PROVIDERS from pr_agent.git_providers.plain_diff_provider import PlainDiffGitProvider +from pr_agent.tools.pr_code_suggestions import PRCodeSuggestions # Diff-mode settings keys these tests mutate on the process-wide singleton. _SETTINGS_KEYS = ["plain_diff.content", "plain_diff.output_path", @@ -173,13 +174,6 @@ def test_temporary_comment_not_emitted(cfg, capsys): assert "Preparing review" not in captured.out -def test_publish_file_comments_not_supported(cfg): - cfg("plain_diff.content", DIFF) - cfg("plain_diff.output_path", None) - provider = PlainDiffGitProvider(None) - assert provider.is_supported("publish_file_comments") is False - - def test_path_traversal_file_not_read(cfg, tmp_path, monkeypatch): # SENTINEL TEST: this test FAILS if the path-traversal guard in # PlainDiffGitProvider.get_diff_files() is removed. @@ -264,9 +258,11 @@ def test_publish_code_suggestions_renders_to_stdout(cfg, capsys): # render the suggestions to stdout. assert provider.publish_code_suggestions(suggestions) is True out = capsys.readouterr().out - assert "Code suggestions" in out - assert "foo.py:2-2" in out - assert "use a constant" in out + assert out == ( + "## Code suggestions\n\n" + "### foo.py:2-2\n" + "**Suggestion:** use a constant\n\n" + ) def test_publish_code_suggestions_empty_is_noop(cfg, capsys): @@ -277,6 +273,61 @@ def test_publish_code_suggestions_empty_is_noop(cfg, capsys): assert capsys.readouterr().out.strip() == "" +@pytest.mark.asyncio +async def test_partial_improve_keeps_suggestions_and_coverage_in_one_output(cfg, capsys, tmp_path): + output = tmp_path / "suggestions.md" + cfg("plain_diff.content", DIFF) + cfg("plain_diff.output_path", str(output)) + provider = PlainDiffGitProvider(None) + assert provider.supports_code_suggestions_artifact() is True + tool = object.__new__(PRCodeSuggestions) + tool.git_provider = provider + tool.failed_chunk_count = 1 + tool.total_chunk_count = 2 + tool.progress_response = None + tool._output_published = False + tool._validate_suggestion = lambda *args: (True, "", True) + + await tool.push_inline_code_suggestions({"code_suggestions": [{ + "relevant_file": "foo.py", + "relevant_lines_start": 2, + "relevant_lines_end": 2, + "suggestion_content": "Use a constant", + "improved_code": "", + "existing_code": "line2-changed", + "label": "maintainability", + "score": 8, + }]}) + + content = output.read_text(encoding="utf-8") + stdout = capsys.readouterr().out + for rendered in (content, stdout): + assert rendered.count("Use a constant") == 1 + assert rendered.count("foo.py:2-2") == 1 + assert rendered.count("1 of 2 analysis chunks failed") == 1 + assert rendered.index("Use a constant") < rendered.index("1 of 2 analysis chunks failed") + + +@pytest.mark.asyncio +async def test_partial_improve_without_suggestions_keeps_coverage_in_artifact(cfg, capsys, tmp_path): + output = tmp_path / "suggestions.md" + cfg("plain_diff.content", DIFF) + cfg("plain_diff.output_path", str(output)) + provider = PlainDiffGitProvider(None) + tool = object.__new__(PRCodeSuggestions) + tool.git_provider = provider + tool.failed_chunk_count = 1 + tool.total_chunk_count = 2 + + await tool.publish_no_suggestions() + + content = output.read_text(encoding="utf-8") + stdout = capsys.readouterr().out + for rendered in (content, stdout): + assert rendered.count("No code suggestions found in the successfully analyzed chunks.") == 1 + assert rendered.count("1 of 2 analysis chunks failed") == 1 + + def test_incremental_review_disabled(cfg): # -i has no meaning for a standalone diff; the provider must disable it so # PRReviewer never takes the incremental path (which would TypeError). diff --git a/tests/unittest/test_pr_add_docs_prompt_contract.py b/tests/unittest/test_pr_add_docs_prompt_contract.py new file mode 100644 index 0000000000..bb63c5c8d9 --- /dev/null +++ b/tests/unittest/test_pr_add_docs_prompt_contract.py @@ -0,0 +1,14 @@ +import yaml + +from pr_agent.config_loader import get_settings + + +def test_add_docs_example_uses_singular_relevant_line(): + prompt = get_settings().pr_add_docs_prompt.system + example_section = prompt.partition("Example output:")[2] + example = example_section.partition("```yaml")[2].partition("```")[0] + + documentation = yaml.safe_load(example)["Code Documentation"][0] + + assert documentation["relevant line"] == 12 + assert "relevant lines" not in documentation diff --git a/tests/unittest/test_pr_code_suggestions_core.py b/tests/unittest/test_pr_code_suggestions_core.py index 6dffc7f9c3..093c846092 100644 --- a/tests/unittest/test_pr_code_suggestions_core.py +++ b/tests/unittest/test_pr_code_suggestions_core.py @@ -914,10 +914,10 @@ async def test_publish_no_suggestions_removes_the_progress_comment_when_quiet(pu git_provider.publish_comment.assert_not_called() -def _provider_with_file(head_file, filename="app.py"): +def _provider_with_file(head_file, filename="app.py", patch=""): git_provider = MagicMock() git_provider.diff_files = [ - FilePatchInfo(base_file="", head_file=head_file, patch="", filename=filename) + FilePatchInfo(base_file="", head_file=head_file, patch=patch, filename=filename) ] git_provider.publish_code_suggestions.return_value = True return git_provider @@ -930,7 +930,11 @@ def _published_suggestion(git_provider): def test_summarized_suggestions_use_the_target_file_indentation(): - git_provider = _provider_with_file("func f() {\n\told()\n}\n", filename="main.go") + git_provider = _provider_with_file( + "func f() {\n\told()\n}\n", + filename="main.go", + patch="@@ -1,3 +1,3 @@\n func f() {\n-\told()\n+\tnew()\n }\n", + ) git_provider.get_line_link.return_value = "https://example.com/main.go#L2" tool = _make_tool(git_provider) suggestion = _valid_suggestion( @@ -980,6 +984,7 @@ def test_summarized_suggestions_normalize_both_sides_of_the_diff(): git_provider = _provider_with_file( "func f() {\n\tif old() {\n\t\tkeep()\n\t}\n}\n", filename="main.go", + patch="@@ -1,5 +1,5 @@\n func f() {\n-\tif old() {\n+\tif new() {\n \t\tkeep()\n \t}\n }\n", ) git_provider.get_line_link.return_value = "https://example.com/main.go#L2-L4" tool = _make_tool(git_provider) diff --git a/tests/unittest/test_pr_code_suggestions_lifecycle.py b/tests/unittest/test_pr_code_suggestions_lifecycle.py index 973c19422d..d3a34646d0 100644 --- a/tests/unittest/test_pr_code_suggestions_lifecycle.py +++ b/tests/unittest/test_pr_code_suggestions_lifecycle.py @@ -4,7 +4,9 @@ import pytest +from pr_agent.algo.types import FilePatchInfo from pr_agent.config_loader import get_settings +from pr_agent.git_providers.plain_diff_provider import PlainDiffGitProvider from pr_agent.tools import pr_code_suggestions as pr_code_suggestions_module from pr_agent.tools.pr_code_suggestions import PRCodeSuggestions from tests.unittest._settings_helpers import restore_settings, snapshot_settings @@ -20,6 +22,40 @@ "github.publish_as_check_run", ) +# A hunk whose new-file side covers lines 1-3, so a suggestion anchored at 1-1 +# passes the diff-hunk range validation in the summarized publish path. +_ANCHORED_HUNK_PATCH = ( + "@@ -1,3 +1,3 @@\n" + " line1\n" + "+line2-changed\n" + " line3\n" +) + + +def _anchored_suggestion(**overrides): + suggestion = { + "score": 1, + "relevant_file": "app.py", + "relevant_lines_start": 1, + "relevant_lines_end": 1, + "label": "maintainability", + "one_sentence_summary": "Anchored suggestion", + } + suggestion.update(overrides) + return suggestion + + +def _provider_with_anchored_diff(provider): + provider.diff_files = [ + FilePatchInfo( + base_file="line1\nline2\nline3\n", + head_file="line1\nline2-changed\nline3\n", + patch=_ANCHORED_HUNK_PATCH, + filename="app.py", + ) + ] + return provider + def _make_tool(provider): # A bare MagicMock returns a truthy mock here, which would thread every publish call. @@ -82,7 +118,7 @@ async def test_run_removes_progress_comment_when_cancelled( async def test_run_does_not_remove_final_summary_when_cancelled_during_dual_publishing(monkeypatch): settings_snapshot = snapshot_settings(_TRACKED_SETTINGS) try: - provider = MagicMock() + provider = _provider_with_anchored_diff(MagicMock()) progress_comment = MagicMock(name="progress_comment") provider.get_files.return_value = [object()] provider.is_supported.return_value = True @@ -95,7 +131,7 @@ async def test_run_does_not_remove_final_summary_when_cancelled_during_dual_publ monkeypatch.setattr( pr_code_suggestions_module, "retry_with_fallback_models", - AsyncMock(return_value={"code_suggestions": [{"score": 1}]}), + AsyncMock(return_value={"code_suggestions": [_anchored_suggestion()]}), ) _configure_published_run() settings = get_settings() @@ -122,7 +158,7 @@ async def test_run_does_not_remove_final_summary_when_cancelled_during_dual_publ async def test_run_does_not_publish_failure_after_successful_summary(monkeypatch, persistent_comment): settings_snapshot = snapshot_settings(_TRACKED_SETTINGS) try: - provider = MagicMock() + provider = _provider_with_anchored_diff(MagicMock()) progress_comment = MagicMock(name="progress_comment") provider.get_files.return_value = [object()] provider.is_supported.return_value = True @@ -135,7 +171,7 @@ async def test_run_does_not_publish_failure_after_successful_summary(monkeypatch monkeypatch.setattr( pr_code_suggestions_module, "retry_with_fallback_models", - AsyncMock(return_value={"code_suggestions": [{"score": 1}]}), + AsyncMock(return_value={"code_suggestions": [_anchored_suggestion()]}), ) _configure_published_run() settings = get_settings() @@ -233,7 +269,9 @@ async def test_run_publishes_failure_when_inline_suggestions_never_publish(monke monkeypatch.setattr( pr_code_suggestions_module, "retry_with_fallback_models", - AsyncMock(return_value={"code_suggestions": [{"score": 1}]}), + AsyncMock(return_value={"code_suggestions": [ + {"score": 1, "relevant_lines_start": 1, "relevant_lines_end": 1} + ]}), ) _configure_published_run() settings = get_settings() @@ -252,7 +290,7 @@ async def test_run_publishes_failure_when_inline_suggestions_never_publish(monke async def test_run_does_not_remove_persistent_summary_when_cancelled_during_dual_publishing(monkeypatch): settings_snapshot = snapshot_settings(_TRACKED_SETTINGS) try: - provider = MagicMock() + provider = _provider_with_anchored_diff(MagicMock()) progress_comment = MagicMock(name="progress_comment") provider.get_files.return_value = [object()] provider.is_supported.return_value = True @@ -267,7 +305,7 @@ async def test_run_does_not_remove_persistent_summary_when_cancelled_during_dual monkeypatch.setattr( pr_code_suggestions_module, "retry_with_fallback_models", - AsyncMock(return_value={"code_suggestions": [{"score": 1}]}), + AsyncMock(return_value={"code_suggestions": [_anchored_suggestion()]}), ) _configure_published_run() settings = get_settings() @@ -322,7 +360,7 @@ async def test_run_preserves_cancellation_when_progress_cleanup_fails(monkeypatc async def test_run_cleans_up_progress_comment_on_check_run_publish(monkeypatch): settings_snapshot = snapshot_settings(_TRACKED_SETTINGS) try: - provider = MagicMock() + provider = _provider_with_anchored_diff(MagicMock()) progress_comment = MagicMock(name="progress_comment") provider.get_files.return_value = [object()] provider.is_supported.return_value = True @@ -336,7 +374,7 @@ async def test_run_cleans_up_progress_comment_on_check_run_publish(monkeypatch): monkeypatch.setattr( pr_code_suggestions_module, "retry_with_fallback_models", - AsyncMock(return_value={"code_suggestions": [{"score": 1}]}), + AsyncMock(return_value={"code_suggestions": [_anchored_suggestion()]}), ) _configure_published_run() settings = get_settings() @@ -360,7 +398,7 @@ async def test_run_cleans_up_progress_comment_on_check_run_publish(monkeypatch): async def test_run_retains_progress_handle_when_check_run_cleanup_fails(monkeypatch): settings_snapshot = snapshot_settings(_TRACKED_SETTINGS) try: - provider = MagicMock() + provider = _provider_with_anchored_diff(MagicMock()) progress_comment = MagicMock(name="progress_comment") provider.get_files.return_value = [object()] provider.is_supported.return_value = True @@ -375,7 +413,7 @@ async def test_run_retains_progress_handle_when_check_run_cleanup_fails(monkeypa monkeypatch.setattr( pr_code_suggestions_module, "retry_with_fallback_models", - AsyncMock(return_value={"code_suggestions": [{"score": 1}]}), + AsyncMock(return_value={"code_suggestions": [_anchored_suggestion()]}), ) _configure_published_run() settings = get_settings() @@ -515,3 +553,179 @@ async def test_failed_inline_retries_preserve_fallback_output( provider.remove_comment.assert_called_once_with(provider.publish_comment.return_value) finally: restore_settings(settings_snapshot) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("publish_output_no_suggestions", [True, False]) +async def test_run_routes_all_invalid_ranges_through_publish_no_suggestions( + monkeypatch, publish_output_no_suggestions +): + settings_snapshot = snapshot_settings( + _TRACKED_SETTINGS + ("pr_code_suggestions.publish_output_no_suggestions",) + ) + try: + provider = MagicMock() + provider.get_files.return_value = [object()] + provider.is_supported.return_value = True + provider.supports_code_suggestions_artifact.return_value = False + tool = _make_tool(provider) + tool.publish_no_suggestions = AsyncMock() + + monkeypatch.setattr( + pr_code_suggestions_module, "retry_with_fallback_models", + AsyncMock(return_value={"code_suggestions": [ + {"relevant_lines_start": -1, "relevant_lines_end": -1, "label": "bug", + "relevant_file": "app.py", "one_sentence_summary": "sentinel"}, + {"relevant_lines_start": 5, "relevant_lines_end": 2, "label": "bug", + "relevant_file": "app.py", "one_sentence_summary": "reversed"}, + ]}), + ) + _configure_published_run() + settings = get_settings() + settings.config.is_auto_command = True + settings.pr_code_suggestions.commitable_code_suggestions = False + settings.pr_code_suggestions.persistent_comment = False + settings.pr_code_suggestions.publish_output_no_suggestions = publish_output_no_suggestions + + await tool.run() + + tool.publish_no_suggestions.assert_awaited_once() + provider.publish_comment.assert_not_called() + assert tool._output_published is False + finally: + restore_settings(settings_snapshot) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("publish_output_no_suggestions", [True, False]) +async def test_run_all_invalid_ranges_honors_quiet_gate_via_real_publish( + monkeypatch, publish_output_no_suggestions +): + settings_snapshot = snapshot_settings( + _TRACKED_SETTINGS + ("pr_code_suggestions.publish_output_no_suggestions",) + ) + try: + provider = MagicMock() + provider.get_files.return_value = [object()] + provider.is_supported.return_value = True + provider.supports_code_suggestions_artifact.return_value = False + provider.publish_comment.return_value = MagicMock() + provider.should_publish_improve_as_thread.return_value = False + tool = _make_tool(provider) + tool.generate_summarized_suggestions = MagicMock(return_value="table") + + monkeypatch.setattr( + pr_code_suggestions_module, "retry_with_fallback_models", + AsyncMock(return_value={"code_suggestions": [ + {"relevant_lines_start": -1, "relevant_lines_end": -1, "label": "bug", + "relevant_file": "app.py", "one_sentence_summary": "sentinel"}, + ]}), + ) + _configure_published_run() + settings = get_settings() + settings.config.is_auto_command = True + settings.pr_code_suggestions.commitable_code_suggestions = False + settings.pr_code_suggestions.persistent_comment = False + settings.pr_code_suggestions.publish_output_no_suggestions = publish_output_no_suggestions + + await tool.run() + + tool.generate_summarized_suggestions.assert_not_called() + if publish_output_no_suggestions: + comments = [call.args[0] for call in provider.publish_comment.call_args_list] + assert len(comments) == 1 + assert "No code suggestions found for the PR." in comments[0] + else: + provider.publish_comment.assert_not_called() + assert get_settings().data["artifact"] == "" + finally: + restore_settings(settings_snapshot) + + +@pytest.mark.asyncio +async def test_run_plain_diff_leaks_no_progress_when_all_ranges_invalid(monkeypatch, capsys): + settings_snapshot = snapshot_settings( + _TRACKED_SETTINGS + ("pr_code_suggestions.publish_output_no_suggestions", "plain_diff.content") + ) + try: + get_settings().set( + "plain_diff.content", + "diff --git a/foo.py b/foo.py\n" + "index 1111111..2222222 100644\n" + "--- a/foo.py\n" + "+++ b/foo.py\n" + "@@ -1,3 +1,3 @@\n" + " line1\n" + "-line2\n" + "+line2-changed\n" + " line3\n", + ) + provider = PlainDiffGitProvider(None) + tool = PRCodeSuggestions.__new__(PRCodeSuggestions) + tool.git_provider = provider + tool.pr_url = "https://example.invalid/pull/1" + tool.progress_response = None + tool.incremental = SimpleNamespace(is_incremental=False) + tool.progress = "## Generating PR code suggestions" + + monkeypatch.setattr( + pr_code_suggestions_module, "retry_with_fallback_models", + AsyncMock(return_value={"code_suggestions": [ + {"relevant_lines_start": -1, "relevant_lines_end": -1, "label": "bug", + "relevant_file": "app.py", "one_sentence_summary": "sentinel"}, + ]}), + ) + _configure_published_run() + settings = get_settings() + settings.config.is_auto_command = False + settings.pr_code_suggestions.commitable_code_suggestions = False + settings.pr_code_suggestions.persistent_comment = False + + await tool.run() + + out = capsys.readouterr().out + # The output-only provider cannot edit/remove a previously published comment, so + # the progress placeholder must never be persisted and only the final document remains. + assert "Generating PR code suggestions" not in out + assert out.count("## PR Code Suggestions") == 1 + assert out.count("No code suggestions found for the PR.") == 1 + finally: + restore_settings(settings_snapshot) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("publish_no_suggestions", [True, False]) +async def test_run_valid_ranges_skip_no_suggestions_comment(monkeypatch, publish_no_suggestions): + settings_snapshot = snapshot_settings(_TRACKED_SETTINGS + ("pr_code_suggestions.publish_output_no_suggestions",)) + try: + provider = MagicMock() + provider.get_files.return_value = [object()] + provider.is_supported.return_value = True + provider.supports_code_suggestions_artifact.return_value = False + provider.publish_comment.return_value = MagicMock() + _provider_with_anchored_diff(provider) + tool = _make_tool(provider) + tool.publish_no_suggestions = AsyncMock() + + monkeypatch.setattr( + pr_code_suggestions_module, "retry_with_fallback_models", + AsyncMock(return_value={"code_suggestions": [ + {"relevant_lines_start": "2", "relevant_lines_end": "2", "label": "bug", + "relevant_file": "app.py", "one_sentence_summary": "valid", + "suggestion_content": "Content", "existing_code": "old()", + "improved_code": "new()", "score": 7}, + ]}), + ) + _configure_published_run() + settings = get_settings() + settings.config.is_auto_command = True + settings.pr_code_suggestions.commitable_code_suggestions = False + settings.pr_code_suggestions.persistent_comment = False + settings.pr_code_suggestions.publish_output_no_suggestions = publish_no_suggestions + + await tool.run() + + tool.publish_no_suggestions.assert_not_awaited() + assert tool._output_published is True + finally: + restore_settings(settings_snapshot) diff --git a/tests/unittest/test_pr_code_suggestions_pr_chat.py b/tests/unittest/test_pr_code_suggestions_pr_chat.py new file mode 100644 index 0000000000..0c141f591c --- /dev/null +++ b/tests/unittest/test_pr_code_suggestions_pr_chat.py @@ -0,0 +1,151 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from pr_agent.algo.types import FilePatchInfo +from pr_agent.config_loader import get_settings +from pr_agent.git_providers.git_provider import GitProvider +from pr_agent.git_providers.github_provider import GithubProvider +from pr_agent.tools import pr_code_suggestions as pr_code_suggestions_module +from pr_agent.tools.pr_code_suggestions import PRCodeSuggestions +from tests.unittest._settings_helpers import restore_settings, snapshot_settings + +_TRACKED_SETTINGS = ( + "config.publish_output", + "config.publish_output_progress", + "config.is_auto_command", + "pr_code_suggestions.commitable_code_suggestions", + "pr_code_suggestions.demand_code_suggestions_self_review", + "pr_code_suggestions.enable_chat_text", + "pr_code_suggestions.enable_help_text", + "pr_code_suggestions.persistent_comment", + "pr_code_suggestions.dual_publishing_score_threshold", +) + + +class _CustomProvider(GitProvider): + """Minimal external provider that inherits the default chat capability.""" + + def is_supported(self, capability: str) -> bool: + return False + + def get_files(self) -> list: + return [] + + def get_diff_files(self) -> list: + return [] + + def publish_description(self, pr_title: str, pr_body: str): + pass + + def publish_code_suggestions(self, code_suggestions: list) -> bool: + return False + + def get_languages(self): + return {} + + def get_pr_branch(self): + return "" + + def get_user_id(self): + return "" + + def get_pr_description_full(self) -> str: + return "" + + def get_repo_settings(self): + return b"" + + def publish_comment(self, pr_comment: str, is_temporary: bool = False): + pass + + def publish_inline_comment(self, body: str, relevant_file: str, relevant_line_in_file: str, original_suggestion=None): + pass + + def publish_inline_comments(self, comments: list[dict]): + pass + + def remove_initial_comment(self): + pass + + def remove_comment(self, comment): + pass + + def get_issue_comments(self): + return [] + + def publish_labels(self, labels): + pass + + def get_pr_labels(self, update=False): + return [] + + def remove_reaction(self, issue_comment_id: int, reaction_id: int) -> bool: + return False + + def get_commit_messages(self) -> str: + return "" + + + +def test_git_provider_supports_pr_chat_defaults_to_false(): + assert _CustomProvider().supports_pr_chat() is False + + +def test_github_provider_supports_pr_chat_is_true(): + assert GithubProvider.supports_pr_chat(MagicMock()) is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("supports_pr_chat", "expect_link"), [(True, True), (False, False)]) +async def test_pr_chat_link_depends_on_provider_capability(monkeypatch, supports_pr_chat, expect_link): + snapshot = snapshot_settings(_TRACKED_SETTINGS) + try: + provider = MagicMock() + provider.get_files.return_value = ["file.py"] + provider.is_supported.side_effect = lambda capability: capability == "gfm_markdown" + provider.supports_pr_chat.return_value = supports_pr_chat + provider.should_publish_improve_as_thread.return_value = False + provider.diff_files = [ + FilePatchInfo( + base_file="old()\n", + head_file="new()\n", + patch="@@ -1,1 +1,1 @@\n-old()\n+new()\n", + filename="file.py", + ) + ] + + tool = PRCodeSuggestions.__new__(PRCodeSuggestions) + tool.pr_url = "https://example.invalid/pr/1" + tool.progress = "progress" + tool.progress_response = None + tool.git_provider = provider + tool.incremental = SimpleNamespace(is_incremental=False) + tool.generate_summarized_suggestions = MagicMock(return_value="## Suggestions") + + async def _fake_retry(*_args, **_kwargs): + return {"code_suggestions": [{"label": "style", "suggestion_content": "clean up", + "relevant_lines_start": 1, "relevant_lines_end": 1, + "relevant_file": "file.py"}]} + + monkeypatch.setattr(pr_code_suggestions_module, "retry_with_fallback_models", _fake_retry) + + settings = get_settings() + settings.config.publish_output = True + settings.config.publish_output_progress = False + settings.config.is_auto_command = True + settings.pr_code_suggestions.commitable_code_suggestions = False + settings.pr_code_suggestions.demand_code_suggestions_self_review = False + settings.pr_code_suggestions.enable_chat_text = True + settings.pr_code_suggestions.enable_help_text = False + settings.pr_code_suggestions.persistent_comment = False + settings.pr_code_suggestions.dual_publishing_score_threshold = 0 + + await tool.run() + + published_body = provider.publish_comment.call_args.args[0] + provider.supports_pr_chat.assert_called_once_with() + assert ("start a [PR chat]" in published_body) is expect_link + finally: + restore_settings(snapshot) diff --git a/tests/unittest/test_pr_code_suggestions_rendering.py b/tests/unittest/test_pr_code_suggestions_rendering.py index 9b0172831f..ac445f087d 100644 --- a/tests/unittest/test_pr_code_suggestions_rendering.py +++ b/tests/unittest/test_pr_code_suggestions_rendering.py @@ -36,6 +36,35 @@ def _suggestion(**overrides): return base +# A hunk whose new-file side covers lines 1-6 of the file. +DEFAULT_PATCH = ( + "@@ -1,5 +1,6 @@\n" + " def f():\n" + "- return old()\n" + "+ return new()\n" + " extra\n" + " more\n" + " lines\n" + "+ added line\n" +) + + +def _diff_file(filename="app.py", patch=DEFAULT_PATCH): + return FilePatchInfo( + base_file="def f():\n return old()\n extra\n more\n lines\n", + head_file="def f():\n return new()\n extra\n more\n lines\n added line\n", + patch=patch, + filename=filename, + ) + + +def _provider_with_diff_files(*filenames): + git_provider = MagicMock() + git_provider.diff_files = [_diff_file(filename) for filename in filenames] + git_provider.get_line_link.return_value = "" + return git_provider + + # --------------------------------------------------------------------------- # _truncate_if_needed # --------------------------------------------------------------------------- @@ -76,6 +105,110 @@ def test_truncate_if_needed_noop_when_under_limit_or_disabled(): restore_settings(snapshot) +@pytest.mark.parametrize( + "suggestion_kwargs", + [ + {"relevant_lines_start": None, "relevant_lines_end": 5}, + {"relevant_lines_start": -1, "relevant_lines_end": -1}, + {"relevant_lines_start": 0, "relevant_lines_end": 0}, + {"relevant_lines_start": -3, "relevant_lines_end": 1}, + {"relevant_lines_start": 5, "relevant_lines_end": 2}, + {"relevant_lines_start": "10", "relevant_lines_end": "bad"}, + {"relevant_lines_start": float("inf"), "relevant_lines_end": 5}, + {"relevant_lines_start": float("-inf"), "relevant_lines_end": 5}, + {"relevant_lines_start": 2, "relevant_lines_end": float("inf")}, + {"relevant_lines_start": float("nan"), "relevant_lines_end": 5}, + {"relevant_lines_start": 2, "relevant_lines_end": float("nan")}, + {"relevant_lines_start": 1.5, "relevant_lines_end": 5}, + {"relevant_lines_start": 2, "relevant_lines_end": 3.5}, + {"relevant_lines_start": True, "relevant_lines_end": True}, + ], +) +def test_is_suggestion_line_range_valid_rejects_unanchorable_ranges(suggestion_kwargs): + tool = _make_tool() + bad = _suggestion(**suggestion_kwargs) + + assert tool._is_suggestion_line_range_valid(bad) is False + + +def test_is_suggestion_line_range_valid_normalizes_valid_range(): + tool = _make_tool(_provider_with_diff_files("app.py")) + good = _suggestion(relevant_lines_start="2", relevant_lines_end="4") + + assert tool._is_suggestion_line_range_valid(good) is True + assert good["relevant_lines_start"] == 2 + assert good["relevant_lines_end"] == 4 + + +def test_is_suggestion_line_range_valid_rejects_missing_keys(): + tool = _make_tool() + suggestion = _suggestion() + suggestion.pop("relevant_lines_start") + suggestion.pop("relevant_lines_end") + + assert tool._is_suggestion_line_range_valid(suggestion) is False + + +def test_is_suggestion_line_range_valid_rejects_file_not_in_diff(): + tool = _make_tool(_provider_with_diff_files("app.py")) + suggestion = _suggestion(relevant_file="other.py") + + assert tool._is_suggestion_line_range_valid(suggestion) is False + + +def test_is_suggestion_line_range_valid_rejects_range_outside_diff(): + tool = _make_tool(_provider_with_diff_files("app.py")) + suggestion = _suggestion(relevant_lines_start=100, relevant_lines_end=100) + + assert tool._is_suggestion_line_range_valid(suggestion) is False + + +def test_is_suggestion_line_range_valid_follows_the_head_file_when_available(): + # Extended context outside the raw hunk is anchorable within the head file. + diff_file = _diff_file() + git_provider = MagicMock() + git_provider.diff_files = [diff_file] + tool = _make_tool(git_provider) + suggestion = _suggestion(relevant_lines_start=5, relevant_lines_end=9) + + assert tool._is_suggestion_line_range_valid(suggestion) is False + + diff_file.head_file += " trailing\n context\n lines\n" + + assert tool._is_suggestion_line_range_valid(suggestion) is True + + +@pytest.mark.parametrize("head_file", ["", "context\n" * 9]) +def test_is_suggestion_line_range_valid_uses_hunk_without_complete_head_file(head_file): + diff_file = _diff_file() + diff_file.head_file = head_file + diff_file.head_file_is_complete = False + git_provider = MagicMock() + git_provider.diff_files = [diff_file] + tool = _make_tool(git_provider) + + assert tool._is_suggestion_line_range_valid( + _suggestion(relevant_lines_start=2, relevant_lines_end=6)) is True + assert tool._is_suggestion_line_range_valid( + _suggestion(relevant_lines_start=5, relevant_lines_end=9)) is False + + +def test_get_patch_range_lines_rejects_oversized_span_without_enumerating_it(): + # A parseable but absurdly large span must be rejected by length comparison + # before any range(start, end) enumeration would iterate over it. + patch = DEFAULT_PATCH + + assert PRCodeSuggestions._get_patch_range_lines(patch, 1, 1_000_000_000) is None + assert PRCodeSuggestions._get_patch_range_lines(patch, 2, 6) == [ + " return new()", + " extra", + " more", + " lines", + " added line", + ] + assert PRCodeSuggestions._get_patch_range_lines(patch, 2, 2) == [" return new()"] + + def test_prepare_pr_code_suggestions_applies_truncation_inline(): settings = get_settings() snapshot = snapshot_settings(TRUNCATION_SETTINGS) @@ -307,7 +440,7 @@ def test_generate_summarized_suggestions_empty_returns_placeholder(): def test_generate_summarized_suggestions_renders_table_and_sorts_by_score(): - git_provider = MagicMock() + git_provider = _provider_with_diff_files("app.py", "auth.py") git_provider.get_line_link.return_value = "https://example.test/app.py#L2" tool = _make_tool(git_provider) settings = get_settings() @@ -343,8 +476,7 @@ def test_generate_summarized_suggestions_renders_table_and_sorts_by_score(): def test_generate_summarized_suggestions_uses_score_string_when_new_mechanism_enabled(): - git_provider = MagicMock() - git_provider.get_line_link.return_value = "" + git_provider = _provider_with_diff_files("app.py") tool = _make_tool(git_provider) settings = get_settings() snapshot = snapshot_settings(["pr_code_suggestions.new_score_mechanism"]) @@ -362,8 +494,7 @@ def test_generate_summarized_suggestions_uses_score_string_when_new_mechanism_en def test_generate_summarized_suggestions_escapes_angle_bracket_strings_in_summary(): - git_provider = MagicMock() - git_provider.get_line_link.return_value = "" + git_provider = _provider_with_diff_files("app.py") tool = _make_tool(git_provider) suggestion = _suggestion(one_sentence_summary="Replace '' with new_name") out = tool.generate_summarized_suggestions({"code_suggestions": [suggestion]}) @@ -375,8 +506,7 @@ def test_generate_summarized_suggestions_escapes_angle_bracket_strings_in_summar def test_generate_summarized_suggestions_includes_score_why_block_when_present(): - git_provider = MagicMock() - git_provider.get_line_link.return_value = "" + git_provider = _provider_with_diff_files("app.py") tool = _make_tool(git_provider) suggestion = _suggestion(score_why="Catches a real bug.") out = tool.generate_summarized_suggestions({"code_suggestions": [suggestion]}) @@ -384,6 +514,136 @@ def test_generate_summarized_suggestions_includes_score_why_block_when_present() assert "Why: Catches a real bug." in out +def test_generate_summarized_suggestions_skips_anchorless_but_keeps_rest(): + """A suggestion without resolved line anchors is skipped instead of failing the whole table.""" + git_provider = _provider_with_diff_files("app.py") + tool = _make_tool(git_provider) + anchored = _suggestion(one_sentence_summary="Keep me") + anchorless = _suggestion(one_sentence_summary="Drop me", relevant_file="other.py") + anchorless.pop("relevant_lines_start") + anchorless.pop("relevant_lines_end") + + out = tool.generate_summarized_suggestions({"code_suggestions": [anchorless, anchored]}) + + assert "" in out + assert "Keep me" in out + assert "Drop me" not in out + + +def test_generate_summarized_suggestions_skips_positive_range_outside_diff(): + """A positive, ordered range outside the file's diff hunks is dropped per-suggestion.""" + git_provider = _provider_with_diff_files("app.py") + tool = _make_tool(git_provider) + in_diff = _suggestion(one_sentence_summary="In diff") + out_diff = _suggestion( + one_sentence_summary="Hallucinated range", + relevant_lines_start=100, + relevant_lines_end=100, + ) + + out = tool.generate_summarized_suggestions({"code_suggestions": [in_diff, out_diff]}) + + assert "
      " in out + assert "In diff" in out + assert "Hallucinated range" not in out + + +def test_generate_summarized_suggestions_all_anchorless_returns_placeholder(): + """When no suggestion has resolvable anchors, publish a truthful message instead of an empty table.""" + git_provider = MagicMock() + git_provider.get_line_link.return_value = "" + tool = _make_tool(git_provider) + suggestion = _suggestion() + suggestion.pop("relevant_lines_start") + suggestion.pop("relevant_lines_end") + + out = tool.generate_summarized_suggestions({"code_suggestions": [suggestion]}) + + assert "No suggestions found to improve this PR." in out + assert "
      " not in out + + +@pytest.mark.parametrize( + "start,end", + [ + (-1, -1), # unresolved -1 sentinel written by self-reflection + (0, 0), + (-3, 1), + (5, 2), # reversed range + ], +) +def test_generate_summarized_suggestions_skips_invalid_line_ranges(start, end): + """Unresolved sentinels, non-positive lines and reversed ranges are omitted per-suggestion.""" + git_provider = _provider_with_diff_files("app.py") + tool = _make_tool(git_provider) + bad = _suggestion(one_sentence_summary="Bad range", relevant_lines_start=start, relevant_lines_end=end) + good = _suggestion(one_sentence_summary="Good range") + + out = tool.generate_summarized_suggestions({"code_suggestions": [bad, good]}) + + assert "
      " in out + assert "Good range" in out + assert "Bad range" not in out + + +@pytest.mark.parametrize("start,end", [(-1, -1), (0, 0), (5, 2)]) +def test_generate_summarized_suggestions_all_invalid_ranges_returns_placeholder(start, end): + git_provider = MagicMock() + git_provider.get_line_link.return_value = "" + tool = _make_tool(git_provider) + bad = _suggestion(relevant_lines_start=start, relevant_lines_end=end) + + out = tool.generate_summarized_suggestions({"code_suggestions": [bad]}) + + assert "No suggestions found to improve this PR." in out + assert "
      " not in out + + +@pytest.mark.asyncio +async def test_analyze_self_reflection_mismatched_count_does_not_crash(): + """Feedback covering a different suggestion count than generated leaves anchors unresolved.""" + tool = _make_tool() + suggestion = _suggestion() + suggestion.pop("relevant_lines_start") + suggestion.pop("relevant_lines_end") + data = {"code_suggestions": [suggestion, _suggestion(one_sentence_summary="Second")]} + + await tool.analyze_self_reflection_response( + data, "code_suggestions:\n- suggestion_score: 8\n why: fast\n" + ) + + assert "relevant_lines_start" not in data["code_suggestions"][0] + assert data["code_suggestions"][1]["relevant_lines_start"] == 2 + + +@pytest.mark.asyncio +async def test_analyze_self_reflection_null_feedback_does_not_crash(): + """A non-null response with `code_suggestions: null` must not abort suggestions.""" + tool = _make_tool() + suggestion = _suggestion() + suggestion.pop("relevant_lines_start") + suggestion.pop("relevant_lines_end") + data = {"code_suggestions": [suggestion]} + + await tool.analyze_self_reflection_response(data, "code_suggestions: null") + + assert "relevant_lines_start" not in data["code_suggestions"][0] + + +@pytest.mark.asyncio +async def test_analyze_self_reflection_non_mapping_feedback_does_not_crash(): + """A top-level list (not a mapping) from reflection must not abort suggestions.""" + tool = _make_tool() + suggestion = _suggestion() + suggestion.pop("relevant_lines_start") + suggestion.pop("relevant_lines_end") + data = {"code_suggestions": [suggestion]} + + await tool.analyze_self_reflection_response(data, "- one\n- two") + + assert "relevant_lines_start" not in data["code_suggestions"][0] + + # --------------------------------------------------------------------------- # Stale one-liner validation # --------------------------------------------------------------------------- diff --git a/tests/unittest/test_pr_command_profiles.py b/tests/unittest/test_pr_command_profiles.py index 1b01003897..e375ae350a 100644 --- a/tests/unittest/test_pr_command_profiles.py +++ b/tests/unittest/test_pr_command_profiles.py @@ -235,8 +235,16 @@ async def test_each_webhook_dispatches_its_default_profile(provider, monkeypatch assert agent.commands == DEFAULT_PR_COMMANDS[provider] -async def test_bitbucket_app_default_profile_passes_the_early_gate(monkeypatch): - calls = [] +async def _run_bitbucket_created_webhook( + monkeypatch, + *, + host_commands, + repo_commands, + eligibility=Eligibility.ELIGIBLE, + repo_should_process=True, + disable_auto_feedback=False, +): + trace = [] payload = { "event": "pullrequest:created", "data": { @@ -253,33 +261,234 @@ async def test_bitbucket_app_default_profile_passes_the_early_gate(monkeypatch): async def get_bearer_token(_shared_secret, _client_key): return "bearer" - async def perform_commands(*args): - calls.append(args) + class RecordingIdentityProvider: + def verify_eligibility(self, *_args): + trace.append("eligibility") + return eligibility + + def should_process_pr_logic(_data): + phase = "repo_filter" if "repo_settings" in trace else "host_filter" + trace.append(phase) + return repo_should_process if phase == "repo_filter" else True + + def apply_repo_settings(_url): + trace.append("repo_settings") + get_settings().set("BITBUCKET_APP.PR_COMMANDS", copy.deepcopy(repo_commands)) + get_settings().set("CONFIG.DISABLE_AUTO_FEEDBACK", disable_auto_feedback) + + async def run_commands(commands, *_args): + trace.append(("dispatch", commands)) - monkeypatch.setattr(server_utils, "get_settings", _Settings) monkeypatch.setattr(bitbucket_app, "is_bot_user", lambda _data: False) - monkeypatch.setattr(bitbucket_app, "should_process_pr_logic", lambda _data: True) + monkeypatch.setattr(bitbucket_app, "should_process_pr_logic", should_process_pr_logic) monkeypatch.setattr(bitbucket_app, "get_fork_safe_secret_provider", lambda: secret_provider) monkeypatch.setattr(bitbucket_app, "get_bearer_token", get_bearer_token) monkeypatch.setattr(bitbucket_app.jwt, "decode", lambda *args, **kwargs: {}) + monkeypatch.setattr(bitbucket_app, "get_identity_provider", RecordingIdentityProvider) + monkeypatch.setattr(bitbucket_app, "apply_repo_settings", apply_repo_settings) + monkeypatch.setattr(bitbucket_app, "_run_commands_bitbucket", run_commands) + endpoint = next(route.endpoint for route in bitbucket_app.router.routes if route.path == "/webhook") + background_tasks = BackgroundTasks() + original_bitbucket_app = copy.deepcopy(global_settings.bitbucket_app) + original_bitbucket = copy.deepcopy(global_settings.get("BITBUCKET", {})) + original_config = copy.deepcopy(global_settings.config) + global_settings.set("BITBUCKET_APP.PR_COMMANDS", copy.deepcopy(host_commands)) + global_settings.set("BITBUCKET.BASE_URL", "https://example.test/app") + global_settings.set("CONFIG.IS_AUTO_COMMAND", False) + + try: + with request_cycle_context({}): + result = await endpoint(background_tasks, _Request(payload)) + await background_tasks() + is_auto_command = get_settings().config.is_auto_command + finally: + global_settings.set("BITBUCKET_APP", original_bitbucket_app) + global_settings.set("BITBUCKET", original_bitbucket) + global_settings.set("CONFIG", original_config) + + return result, trace, is_auto_command + + +async def test_bitbucket_app_created_event_loads_repo_commands_before_dispatch(monkeypatch): + result, trace, is_auto_command = await _run_bitbucket_created_webhook( + monkeypatch, + host_commands=[], + repo_commands=["/review"], + ) + + assert result == "OK" + assert trace == [ + "host_filter", + "eligibility", + "repo_settings", + "repo_filter", + ("dispatch", ["/review"]), + ] + assert is_auto_command is True + + +async def test_bitbucket_app_created_event_skips_empty_repo_commands(monkeypatch): + result, trace, is_auto_command = await _run_bitbucket_created_webhook( + monkeypatch, + host_commands=["/review"], + repo_commands=[], + ) + + assert result == "OK" + assert trace == ["host_filter", "eligibility", "repo_settings", "repo_filter"] + assert is_auto_command is False + + +async def test_bitbucket_app_created_event_applies_repo_filter_before_dispatch(monkeypatch): + result, trace, is_auto_command = await _run_bitbucket_created_webhook( + monkeypatch, + host_commands=["/review"], + repo_commands=["/review"], + repo_should_process=False, + ) + + assert result == "OK" + assert trace == ["host_filter", "eligibility", "repo_settings", "repo_filter"] + assert is_auto_command is False + + +async def test_bitbucket_app_created_event_honors_repo_disable_auto_feedback(monkeypatch): + result, trace, is_auto_command = await _run_bitbucket_created_webhook( + monkeypatch, + host_commands=["/review"], + repo_commands=["/review"], + disable_auto_feedback=True, + ) + + assert result == "OK" + assert trace == ["host_filter", "eligibility", "repo_settings"] + assert is_auto_command is False + + +async def test_bitbucket_app_created_event_does_not_load_repo_settings_when_ineligible(monkeypatch): + result, trace, is_auto_command = await _run_bitbucket_created_webhook( + monkeypatch, + host_commands=[], + repo_commands=["/review"], + eligibility=Eligibility.NOT_ELIGIBLE, + ) + + assert result == "OK" + assert trace == ["host_filter", "eligibility"] + assert is_auto_command is False + + +async def _run_bitbucket_push_webhook( + monkeypatch, + *, + host_trigger, + host_commands, + repo_trigger, + repo_commands, +): + calls = [] + pr_url = "https://example.test/pr/1" + payload = { + "event": "pullrequest:updated", + "data": { + "actor": {"type": "user", "account_id": "account"}, + "pullrequest": {"links": {"html": {"href": pr_url}}}, + }, + } + secret_provider = type( + "SecretProvider", + (), + {"get_secret": lambda self, _key: json.dumps({"shared_secret": "secret"})}, + )() + + async def get_bearer_token(_shared_secret, _client_key): + return "bearer" + + def apply_repo_settings(_url): + calls.append("settings") + get_settings().set("BITBUCKET_APP.HANDLE_PUSH_TRIGGER", repo_trigger) + get_settings().set("BITBUCKET_APP.PUSH_COMMANDS", list(repo_commands)) + + async def validate_push(_data): + calls.append("validate") + return True + + @asynccontextmanager + async def record_slot(_key, **_kwargs): + calls.append("slot") + yield True + + async def run_commands(commands, *_args): + calls.append(("commands", list(commands))) + + monkeypatch.setattr(bitbucket_app, "is_bot_user", lambda _data: False) + monkeypatch.setattr(bitbucket_app, "get_fork_safe_secret_provider", lambda: secret_provider) + monkeypatch.setattr(bitbucket_app, "get_bearer_token", get_bearer_token) + monkeypatch.setattr(bitbucket_app.jwt, "decode", lambda *args, **kwargs: {}) + monkeypatch.setattr(bitbucket_app, "get_identity_provider", _IdentityProvider) + monkeypatch.setattr(bitbucket_app, "apply_repo_settings", apply_repo_settings) monkeypatch.setattr( bitbucket_app, - "get_identity_provider", - _IdentityProvider, + "should_process_pr_logic", + lambda _data: pytest.fail("Bitbucket update events must not use PR creation filters"), ) - monkeypatch.setattr(bitbucket_app, "_perform_commands_bitbucket", perform_commands) + monkeypatch.setattr(bitbucket_app, "_validate_time_from_last_commit_to_pr_update", validate_push) + monkeypatch.setattr(bitbucket_app, "push_trigger_slot", record_slot) + monkeypatch.setattr(bitbucket_app, "_run_commands_bitbucket", run_commands) endpoint = next(route.endpoint for route in bitbucket_app.router.routes if route.path == "/webhook") background_tasks = BackgroundTasks() + original_push_commands = global_settings.get("BITBUCKET_APP.PUSH_COMMANDS", None) + original_handle_push = global_settings.get("BITBUCKET_APP.HANDLE_PUSH_TRIGGER", None) + original_base_url = global_settings.get("BITBUCKET.BASE_URL", None) + global_settings.set("BITBUCKET_APP.HANDLE_PUSH_TRIGGER", host_trigger) + global_settings.set("BITBUCKET_APP.PUSH_COMMANDS", list(host_commands)) + global_settings.set("BITBUCKET.BASE_URL", "https://example.test/app") - with request_cycle_context({}): - context["settings"] = copy.deepcopy(global_settings) - get_settings().set("BITBUCKET.BASE_URL", "https://example.test/app") - result = await endpoint(background_tasks, _Request(payload)) - await background_tasks() + try: + with request_cycle_context({}): + result = await endpoint(background_tasks, _Request(payload)) + await background_tasks() + finally: + global_settings.set("BITBUCKET_APP.PUSH_COMMANDS", original_push_commands) + global_settings.set("BITBUCKET_APP.HANDLE_PUSH_TRIGGER", original_handle_push) + global_settings.set("BITBUCKET.BASE_URL", original_base_url) + + return result, calls + + +async def test_bitbucket_app_push_applies_repo_settings_before_effective_gate(monkeypatch): + result, calls = await _run_bitbucket_push_webhook( + monkeypatch, + host_trigger=False, + host_commands=[], + repo_trigger=True, + repo_commands=["/review"], + ) + + assert result == "OK" + assert calls == ["settings", "validate", "slot", ("commands", ["/review"])] + + +@pytest.mark.parametrize( + ("repo_trigger", "repo_commands"), + [ + (False, ["/review"]), + (True, []), + ], +) +async def test_bitbucket_app_effective_push_config_skips_before_validation( + monkeypatch, repo_trigger, repo_commands +): + result, calls = await _run_bitbucket_push_webhook( + monkeypatch, + host_trigger=True, + host_commands=["/host-review"], + repo_trigger=repo_trigger, + repo_commands=repo_commands, + ) assert result == "OK" - assert len(calls) == 1 - assert calls[0][0] == "pr_commands" + assert calls == ["settings"] @pytest.mark.parametrize("proceed", [True, False]) diff --git a/tests/unittest/test_pr_description.py b/tests/unittest/test_pr_description.py index fb2d2ae27a..612d3e6218 100644 --- a/tests/unittest/test_pr_description.py +++ b/tests/unittest/test_pr_description.py @@ -1,11 +1,13 @@ +import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest import yaml from jinja2 import Environment, StrictUndefined +from pr_agent.algo.pr_processing import retry_with_fallback_models from pr_agent.algo.types import FilePatchInfo -from pr_agent.algo.utils import load_yaml +from pr_agent.algo.utils import ModelType, load_yaml from pr_agent.config_loader import get_settings from pr_agent.tools.pr_description import ( PRDescription, @@ -68,6 +70,29 @@ def _make_large_pr_instance(diff_files=None): return obj +def _large_pr_chunks(): + chunks = [ + ["diff --git a/src/file1.py b/src/file1.py\n... file1 ..."], + ["diff --git a/src/file2.py b/src/file2.py\n... file2 ..."], + ] + return chunks, [10, 10], [], [], {}, [["src/file1.py"], ["src/file2.py"]] + + +def _file_prediction(filename: str) -> str: + return f"""pr_files: +- filename: {filename} + changes_title: Describe {filename} + changes_summary: File summary + label: enhancement""" + + +def _header_prediction() -> str: + return """type: +- Enhancement +title: Partial description +description: Summarizes the successfully described files.""" + + def _mock_settings(pr_diagram_direction: str = 'adaptive', pr_diagram_direction_threshold: int = 5): """Mock get_settings used by _prepare_data.""" settings = MagicMock() @@ -276,14 +301,12 @@ def test_prepare_pr_answer_with_markers_replaces_plain_and_comment_markers(self, "changes_diagram": "\n```mermaid\ngraph LR\nA --> B\n```", } - title, body, walkthrough, file_changes = obj._prepare_pr_answer_with_markers() + title, body = obj._prepare_pr_answer_with_markers() assert title == "AI title" assert "Bug fix" in body assert "Fixes the cache invalidation bug." in body assert "```mermaid" in body - assert walkthrough == "" - assert file_changes == [] @pytest.mark.asyncio async def test_extend_uncovered_files_adds_missing_diff_files_to_prediction(self): @@ -583,6 +606,235 @@ async def mock_get_prediction(model, patches_diff, prompt="pr_description_prompt assert obj.data["title"].strip() == "Combined PR Title" assert [f["filename"].strip() for f in obj.data["pr_files"]] == ["src/file1.py", "src/file2.py"] + @pytest.mark.asyncio + @pytest.mark.parametrize("async_calls", [True, False]) + async def test_large_pr_retains_valid_chunk_when_sibling_raises(self, monkeypatch, async_calls): + obj = _make_large_pr_instance() + monkeypatch.setattr(get_settings().pr_description, "async_ai_calls", async_calls) + + async def mock_get_prediction(model, patches_diff, prompt="pr_description_prompt"): + if prompt == "pr_description_only_description_prompts": + return _header_prediction() + if "file1" in patches_diff: + return _file_prediction("src/file1.py") + raise RuntimeError("file2 transport failed") + + obj._get_prediction = AsyncMock(side_effect=mock_get_prediction) + with patch("pr_agent.tools.pr_description.get_pr_diff", return_value=""), patch( + "pr_agent.tools.pr_description.get_pr_diff_multiple_patchs", + return_value=_large_pr_chunks(), + ): + await obj._prepare_prediction("gpt-4o") + + parsed = load_yaml(obj.prediction, keys_fix_yaml=obj.keys_fix) + assert [file["filename"].strip() for file in parsed["pr_files"]] == ["src/file1.py", "src/file2.py"] + assert parsed["pr_files"][1]["label"].strip() == "additional files" + assert obj.description_failed_chunk_count == 1 + assert obj.description_total_chunk_count == 2 + assert obj.description_failed_files == ["src/file2.py"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("malformed_prediction", [ + "pr_files: []", + "pr_files:\n- not-a-file-record", + "pr_files:\n- {}", + """pr_files: +- filename: src/file2.py + changes_title: Missing label + changes_summary: File summary + label:""", + ]) + async def test_large_pr_retains_valid_chunk_when_sibling_is_malformed( + self, monkeypatch, malformed_prediction): + obj = _make_large_pr_instance() + monkeypatch.setattr(get_settings().pr_description, "async_ai_calls", True) + + async def mock_get_prediction(model, patches_diff, prompt="pr_description_prompt"): + if prompt == "pr_description_only_description_prompts": + return _header_prediction() + if "file1" in patches_diff: + return _file_prediction("src/file1.py") + return malformed_prediction + + obj._get_prediction = AsyncMock(side_effect=mock_get_prediction) + with patch("pr_agent.tools.pr_description.get_pr_diff", return_value=""), patch( + "pr_agent.tools.pr_description.get_pr_diff_multiple_patchs", + return_value=_large_pr_chunks(), + ): + await obj._prepare_prediction("gpt-4o") + + parsed = load_yaml(obj.prediction, keys_fix_yaml=obj.keys_fix) + assert [file["filename"].strip() for file in parsed["pr_files"]] == ["src/file1.py", "src/file2.py"] + assert obj.description_failed_chunk_count == 1 + assert obj.description_failed_files == ["src/file2.py"] + + @pytest.mark.asyncio + async def test_large_pr_keeps_complete_records_from_mixed_chunk(self, monkeypatch): + obj = _make_large_pr_instance() + monkeypatch.setattr(get_settings().pr_description, "async_ai_calls", True) + + async def mock_get_prediction(model, patches_diff, prompt="pr_description_prompt"): + if prompt == "pr_description_only_description_prompts": + return _header_prediction() + return """pr_files: +- filename: src/file1.py + changes_title: Describe src/file1.py + changes_summary: File summary + label: enhancement +- filename: src/file2.py + changes_title: Missing summary + label: enhancement""" + + obj._get_prediction = AsyncMock(side_effect=mock_get_prediction) + with patch("pr_agent.tools.pr_description.get_pr_diff", return_value=""), patch( + "pr_agent.tools.pr_description.get_pr_diff_multiple_patchs", + return_value=( + [["diff --git a/src/file1.py b/src/file1.py\n... file1 ...", + "diff --git a/src/file2.py b/src/file2.py\n... file2 ..."]], + [20], + [], + [], + {}, + [["src/file1.py", "src/file2.py"]], + ), + ): + await obj._prepare_prediction("gpt-4o") + + parsed = load_yaml(obj.prediction, keys_fix_yaml=obj.keys_fix) + assert [file["filename"].strip() for file in parsed["pr_files"]] == ["src/file1.py", "src/file2.py"] + obj._prepare_data() + assert obj._prepare_file_labels() == { + "enhancement": [("src/file1.py", "Describe src/file1.py", "File summary")] + } + + @pytest.mark.asyncio + @pytest.mark.parametrize("async_calls", [True, False]) + async def test_large_pr_all_failed_chunks_raise_for_model_fallback(self, monkeypatch, async_calls): + obj = _make_large_pr_instance() + monkeypatch.setattr(get_settings().pr_description, "async_ai_calls", async_calls) + chunk_calls = [] + + async def mock_get_prediction(model, patches_diff, prompt="pr_description_prompt"): + chunk_calls.append(patches_diff) + raise RuntimeError(f"{patches_diff.splitlines()[0]} failed") + + obj._get_prediction = AsyncMock(side_effect=mock_get_prediction) + with patch("pr_agent.tools.pr_description.get_pr_diff", return_value=""), patch( + "pr_agent.tools.pr_description.get_pr_diff_multiple_patchs", + return_value=_large_pr_chunks(), + ), pytest.raises(RuntimeError, match="file1.py"): + await obj._prepare_prediction("gpt-4o") + + assert len(chunk_calls) == 2 + + @pytest.mark.asyncio + async def test_large_pr_all_malformed_chunks_raise_for_model_fallback(self, monkeypatch): + obj = _make_large_pr_instance() + monkeypatch.setattr(get_settings().pr_description, "async_ai_calls", True) + + obj._get_prediction = AsyncMock(return_value="pr_files:\n- not-a-file-record") + with patch("pr_agent.tools.pr_description.get_pr_diff", return_value=""), patch( + "pr_agent.tools.pr_description.get_pr_diff_multiple_patchs", + return_value=_large_pr_chunks(), + ), pytest.raises(ValueError, match="invalid YAML"): + await obj._prepare_prediction("gpt-4o") + + assert all( + call.kwargs["prompt"] == "pr_description_only_files_prompts" + for call in obj._get_prediction.await_args_list + ) + + @pytest.mark.asyncio + async def test_large_pr_empty_chunk_keeps_result_associated_with_its_files(self, monkeypatch): + obj = _make_large_pr_instance() + monkeypatch.setattr(get_settings().pr_description, "async_ai_calls", True) + chunks = [[], ["diff --git a/src/file2.py b/src/file2.py\n... file2 ..."]] + + async def mock_get_prediction(model, patches_diff, prompt="pr_description_prompt"): + if prompt == "pr_description_only_description_prompts": + return _header_prediction() + return _file_prediction("src/file2.py") + + obj._get_prediction = AsyncMock(side_effect=mock_get_prediction) + with patch("pr_agent.tools.pr_description.get_pr_diff", return_value=""), patch( + "pr_agent.tools.pr_description.get_pr_diff_multiple_patchs", + return_value=(chunks, [0, 10], [], [], {}, [["src/file1.py"], ["src/file2.py"]]), + ): + await obj._prepare_prediction("gpt-4o") + + assert obj.description_failed_files == ["src/file1.py"] + assert obj.description_failed_chunk_count == 1 + assert len(obj._get_prediction.await_args_list) == 2 # one non-empty chunk and the header pass + + @pytest.mark.asyncio + async def test_large_pr_chunk_cancellation_propagates(self, monkeypatch): + obj = _make_large_pr_instance() + monkeypatch.setattr(get_settings().pr_description, "async_ai_calls", True) + + async def mock_get_prediction(model, patches_diff, prompt="pr_description_prompt"): + if "file2" in patches_diff: + raise asyncio.CancelledError() + return _file_prediction("src/file1.py") + + obj._get_prediction = AsyncMock(side_effect=mock_get_prediction) + with patch("pr_agent.tools.pr_description.get_pr_diff", return_value=""), patch( + "pr_agent.tools.pr_description.get_pr_diff_multiple_patchs", + return_value=_large_pr_chunks(), + ), pytest.raises(asyncio.CancelledError): + await obj._prepare_prediction("gpt-4o") + + @pytest.mark.asyncio + async def test_large_pr_all_fail_then_fallback_model_keeps_partial_result(self, monkeypatch): + obj = _make_large_pr_instance() + monkeypatch.setattr(get_settings().pr_description, "async_ai_calls", True) + calls = [] + + async def mock_get_prediction(model, patches_diff, prompt="pr_description_prompt"): + calls.append((model, prompt, patches_diff)) + if model == "gpt-4o": + raise RuntimeError("primary chunk failed") + if prompt == "pr_description_only_description_prompts": + return _header_prediction() + if "file1" in patches_diff: + return _file_prediction("src/file1.py") + raise RuntimeError("fallback file2 failed") + + obj._get_prediction = AsyncMock(side_effect=mock_get_prediction) + with patch("pr_agent.tools.pr_description.get_pr_diff", return_value=""), patch( + "pr_agent.tools.pr_description.get_pr_diff_multiple_patchs", + return_value=_large_pr_chunks(), + ), patch( + "pr_agent.algo.pr_processing._get_all_models", return_value=["gpt-4o", "gpt-4o-mini"] + ), patch( + "pr_agent.algo.pr_processing._get_all_deployments", return_value=[None, None] + ), patch( + "pr_agent.algo.pr_processing.route_primary_model", return_value=None + ): + await retry_with_fallback_models(obj._prepare_prediction, ModelType.WEAK) + + assert [(model, prompt) for model, prompt, _ in calls] == [ + ("gpt-4o", "pr_description_only_files_prompts"), + ("gpt-4o", "pr_description_only_files_prompts"), + ("gpt-4o-mini", "pr_description_only_files_prompts"), + ("gpt-4o-mini", "pr_description_only_files_prompts"), + ("gpt-4o-mini", "pr_description_only_description_prompts"), + ] + assert obj.description_failed_chunk_count == 1 + assert obj.description_failed_files == ["src/file2.py"] + + def test_description_coverage_footer_names_failed_chunks_and_files(self): + obj = _make_large_pr_instance() + obj.description_failed_chunk_count = 1 + obj.description_total_chunk_count = 2 + obj.description_failed_files = ["src/file2.py"] + + footer = obj._get_description_coverage_footer() + + assert "1 of 2 file-description chunks failed" in footer + assert "based on the successful chunks only" in footer + assert "`src/file2.py`" in footer + assert "
      " not in footer + @pytest.mark.asyncio async def test_prepare_prediction_normal_diff_uses_single_prompt(self): """Verify normal diff under token limit does NOT fall into the large-PR multi-patch path.""" @@ -632,6 +884,7 @@ async def mock_get_prediction(model, patches_diff, prompt="pr_description_prompt return """pr_files: - filename: src/file1.py changes_title: Add diagram support + changes_summary: Adds diagram support. label: enhancement""" elif prompt == "pr_description_only_description_prompts": return """type: diff --git a/tests/unittest/test_pr_description_help_rendering.py b/tests/unittest/test_pr_description_help_rendering.py new file mode 100644 index 0000000000..110459312f --- /dev/null +++ b/tests/unittest/test_pr_description_help_rendering.py @@ -0,0 +1,116 @@ +"""The /describe help footer picks list markup from provider capabilities.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from pr_agent.config_loader import get_settings +from pr_agent.git_providers.github_provider import GithubProvider +from pr_agent.git_providers.gitlab_provider import GitLabProvider +from pr_agent.tools import pr_description as pr_description_module +from pr_agent.tools.pr_description import PRDescription +from tests.unittest._settings_helpers import restore_settings, snapshot_settings + +_TRACKED_SETTINGS = ( + "config.is_auto_command", + "config.propagate_tool_errors", + "config.publish_output", + "pr_description.enable_help_comment", + "pr_description.enable_help_text", + "pr_description.enable_semantic_files_types", + "pr_description.final_update_message", + "pr_description.generate_ai_title", + "pr_description.publish_description_as_comment", + "pr_description.publish_labels", + "pr_description.use_description_markers", +) + +HTML_LIST_MARKER = "
    • Type /help how to ..." +PLAIN_BULLET_MARKER = "- Type /help how to ..." +PLAIN_BREAK_MARKER = "PR-Agent usage.
      - Check out the" + + +class StubProvider: + """A provider that is none of the concrete classes the tool used to branch on.""" + + def __init__(self, html_lists: bool): + self._html_lists = html_lists + self.html_lists_queries = 0 + self.descriptions = [] + + def is_supported(self, capability: str) -> bool: + return capability == "gfm_markdown" + + def supports_inline_help_footer(self) -> bool: + self.html_lists_queries += 1 + return self._html_lists + + def publish_description(self, title: str, body: str): + self.descriptions.append((title, body)) + + +@pytest.fixture +def configured_settings(): + snapshot = snapshot_settings(_TRACKED_SETTINGS) + settings = get_settings() + settings.config.is_auto_command = True + settings.config.propagate_tool_errors = True + settings.config.publish_output = True + settings.pr_description.enable_help_comment = True + settings.pr_description.enable_help_text = False + settings.pr_description.enable_semantic_files_types = False + settings.pr_description.final_update_message = False + settings.pr_description.generate_ai_title = True + settings.pr_description.publish_description_as_comment = False + settings.pr_description.publish_labels = False + settings.pr_description.use_description_markers = False + yield + restore_settings(snapshot) + + +async def render_description_help(provider, monkeypatch) -> str: + description = PRDescription.__new__(PRDescription) + description.pr_id = "1" + description.git_provider = provider + description.vars = {} + description.prediction = "generated" + description.data = None + description.file_label_dict = None + description._prepare_data = MagicMock() + description._prepare_pr_answer = MagicMock(return_value=("AI title", "Description", "")) + + monkeypatch.setattr(pr_description_module, "extract_and_cache_pr_tickets", AsyncMock()) + monkeypatch.setattr(pr_description_module, "retry_with_fallback_models", AsyncMock()) + + await description.run() + + assert len(provider.descriptions) == 1 + return provider.descriptions[0][1] + + +@pytest.mark.parametrize( + "html_lists, expected, unexpected", + [ + (True, HTML_LIST_MARKER, PLAIN_BULLET_MARKER), + (False, PLAIN_BULLET_MARKER, HTML_LIST_MARKER), + ], +) +async def test_description_help_rendering_follows_html_list_capability( + configured_settings, monkeypatch, html_lists, expected, unexpected +): + provider = StubProvider(html_lists) + + body = await render_description_help(provider, monkeypatch) + + assert provider.html_lists_queries == 1 + assert expected in body + assert unexpected not in body + assert "https://docs.pr-agent.ai/usage-guide/" in body + assert "qodo-merge-docs.qodo.ai" not in body + if not html_lists: + assert PLAIN_BREAK_MARKER in body + + +def test_providers_declare_html_list_capability(): + assert GithubProvider.__new__(GithubProvider).supports_inline_help_footer() is True + assert GitLabProvider.__new__(GitLabProvider).supports_inline_help_footer() is False diff --git a/tests/unittest/test_pr_description_lifecycle.py b/tests/unittest/test_pr_description_lifecycle.py index a00f0113ac..41f5a4f433 100644 --- a/tests/unittest/test_pr_description_lifecycle.py +++ b/tests/unittest/test_pr_description_lifecycle.py @@ -148,6 +148,40 @@ async def test_run_neutralizes_progress_comment_before_delete(monkeypatch): restore_settings(settings_snapshot) +@pytest.mark.asyncio +async def test_run_appends_partial_description_coverage_to_output(monkeypatch): + settings_snapshot = snapshot_settings(_TRACKED_SETTINGS) + try: + provider = MagicMock() + provider.is_supported.return_value = False + description = _make_description(provider) + description.prediction = "generated" + description.description_failed_chunk_count = 1 + description.description_total_chunk_count = 2 + description.description_failed_files = ["src/file2.py"] + description._prepare_data = MagicMock() + description._prepare_pr_answer = MagicMock(return_value=("AI title", "Description", "")) + + monkeypatch.setattr(pr_description_module, "extract_and_cache_pr_tickets", AsyncMock()) + monkeypatch.setattr(pr_description_module, "retry_with_fallback_models", AsyncMock()) + settings = get_settings() + settings.config.publish_output = False + settings.pr_description.enable_help_comment = False + settings.pr_description.enable_help_text = False + settings.pr_description.enable_semantic_files_types = False + settings.pr_description.publish_labels = False + settings.pr_description.use_description_markers = False + + await description.run() + + artifact = get_settings().data["artifact"] + assert "Description coverage" in artifact + assert "1 of 2 file-description chunks failed" in artifact + assert "`src/file2.py`" in artifact + finally: + restore_settings(settings_snapshot) + + @pytest.mark.asyncio @pytest.mark.parametrize("propagate_tool_errors", [False, True]) async def test_run_reports_description_publication_failure(monkeypatch, propagate_tool_errors): @@ -161,7 +195,7 @@ async def test_run_reports_description_publication_failure(monkeypatch, propagat description = _make_description(provider) description.prediction = "generated" description._prepare_data = MagicMock() - description._prepare_pr_answer = MagicMock(return_value=("AI title", "Description", "", [])) + description._prepare_pr_answer = MagicMock(return_value=("AI title", "Description", "")) monkeypatch.setattr(pr_description_module, "extract_and_cache_pr_tickets", AsyncMock()) monkeypatch.setattr(pr_description_module, "retry_with_fallback_models", AsyncMock()) @@ -189,3 +223,40 @@ async def test_run_reports_description_publication_failure(monkeypatch, propagat provider.remove_comment.assert_called_once_with(progress_comment) finally: restore_settings(settings_snapshot) + + +@pytest.mark.asyncio +async def test_run_always_includes_walkthrough_in_the_description_body(monkeypatch): + # Regression guard for #3116: a provider advertising inline-comment capability support used + # to strip the File Walkthrough from the description body and publish nothing in its place. + # The walkthrough must always be appended on the non-marker rendering path. + settings_snapshot = snapshot_settings(_TRACKED_SETTINGS) + try: + provider = MagicMock() + provider.is_supported.return_value = True + description = _make_description(provider) + description.prediction = "generated" + description._prepare_data = MagicMock() + description._prepare_pr_answer = MagicMock(return_value=("AI title", "Base body", "WALKTHROUGH")) + + monkeypatch.setattr(pr_description_module, "extract_and_cache_pr_tickets", AsyncMock()) + monkeypatch.setattr(pr_description_module, "retry_with_fallback_models", AsyncMock()) + _configure_published_run() + settings = get_settings() + settings.pr_description.enable_help_comment = False + settings.pr_description.enable_help_text = False + settings.pr_description.enable_semantic_files_types = False + settings.pr_description.final_update_message = False + settings.pr_description.generate_ai_title = True + settings.pr_description.publish_description_as_comment = False + settings.pr_description.publish_labels = False + settings.pr_description.use_description_markers = False + + await description.run() + + provider.publish_description.assert_called_once() + title, body = provider.publish_description.call_args.args + assert title == "AI title" + assert "WALKTHROUGH" in body + finally: + restore_settings(settings_snapshot) diff --git a/tests/unittest/test_pr_description_output_core.py b/tests/unittest/test_pr_description_output_core.py index 9d7a6a6352..421979932f 100644 --- a/tests/unittest/test_pr_description_output_core.py +++ b/tests/unittest/test_pr_description_output_core.py @@ -190,7 +190,7 @@ def test_html_comment_guard_prevents_type_replacement(self, mock_get_settings): body_in = "\npr_agent:type stays raw" obj = self._obj_with_user_description(body_in, {"title": "AI", "type": "Bug fix"}) - _, body, _, _ = obj._prepare_pr_answer_with_markers() + _, body = obj._prepare_pr_answer_with_markers() # Guard present -> the plain marker is NOT replaced. assert "pr_agent:type stays raw" in body @@ -204,7 +204,7 @@ def test_plain_summary_marker_is_replaced(self, mock_get_settings): {"title": "AI", "description": "Adds caching layer."}, ) - _, body, _, _ = obj._prepare_pr_answer_with_markers() + _, body = obj._prepare_pr_answer_with_markers() assert "Adds caching layer." in body assert "pr_agent:summary" not in body @@ -217,7 +217,7 @@ def test_summary_marker_is_removed_when_description_disabled(self, mock_get_sett {"title": "AI", "description": "Adds caching layer."}, ) - _, body, _, _ = obj._prepare_pr_answer_with_markers() + _, body = obj._prepare_pr_answer_with_markers() assert "Adds caching layer." not in body assert "pr_agent:summary" not in body @@ -230,7 +230,7 @@ def test_html_comment_summary_marker_is_removed_when_description_disabled(self, {"title": "AI", "description": "Adds caching layer."}, ) - _, body, _, _ = obj._prepare_pr_answer_with_markers() + _, body = obj._prepare_pr_answer_with_markers() assert "Adds caching layer." not in body assert "pr_agent:summary" not in body @@ -243,7 +243,7 @@ def test_generated_by_header_prefixes_replacements(self, mock_get_settings): {"title": "AI", "type": "Bug fix", "description": "Fix bug."}, ) - _, body, _, _ = obj._prepare_pr_answer_with_markers() + _, body = obj._prepare_pr_answer_with_markers() assert "### 🤖 Generated by PR Agent at deadbeef" in body # Header appears for both replaced markers. @@ -257,7 +257,7 @@ def test_list_type_is_joined_with_comma(self, mock_get_settings): {"title": "AI", "type": ["Bug fix", "Refactor"]}, ) - _, body, _, _ = obj._prepare_pr_answer_with_markers() + _, body = obj._prepare_pr_answer_with_markers() assert "Bug fix, Refactor" in body @@ -270,7 +270,7 @@ def test_diagram_marker_replaces_both_plain_and_html_comment(self, mock_get_sett {"title": "AI", "changes_diagram": diagram}, ) - _, body, _, _ = obj._prepare_pr_answer_with_markers() + _, body = obj._prepare_pr_answer_with_markers() # Both forms are substituted with the diagram. assert body.count("```mermaid") == 2 @@ -285,7 +285,7 @@ def test_title_falls_back_when_generate_ai_title_disabled(self, mock_get_setting {"title": "AI Title", "description": "x"}, ) - title, _, _, _ = obj._prepare_pr_answer_with_markers() + title, _ = obj._prepare_pr_answer_with_markers() assert title == "Original title" @@ -314,7 +314,7 @@ def test_labels_removed_when_provider_supports_get_labels(self, mock_get_setting obj = self._obj({"title": "t", "labels": ["bug"], "description": "d"}) obj.git_provider.is_supported.side_effect = lambda cap: cap in {"gfm_markdown", "get_labels"} - _, body, _, _ = obj._prepare_pr_answer() + _, body, _ = obj._prepare_pr_answer() # The Labels section is suppressed for providers with native label support. assert "Labels" not in body @@ -325,7 +325,7 @@ def test_type_section_removed_when_disabled(self, mock_get_settings): mock_get_settings.return_value = _settings(enable_pr_type=False) obj = self._obj({"title": "t", "type": "Bug fix", "description": "d"}) - _, body, _, _ = obj._prepare_pr_answer() + _, body, _ = obj._prepare_pr_answer() assert "PR Type" not in body assert "Bug fix" not in body @@ -335,7 +335,7 @@ def test_description_section_removed_when_disabled(self, mock_get_settings): mock_get_settings.return_value = _settings(enable_pr_description=False) obj = self._obj({"title": "t", "type": "Bug fix", "description": "AI summary"}) - _, body, _, _ = obj._prepare_pr_answer() + _, body, _ = obj._prepare_pr_answer() assert "AI summary" not in body # The other sections are untouched. @@ -353,7 +353,7 @@ def test_diagram_and_walkthrough_survive_a_disabled_description(self, mock_get_s "pr_files": [{"filename": "app.py", "changes_title": "c", "label": "enhancement"}], }) - _, body, walkthrough, _ = obj._prepare_pr_answer() + _, body, walkthrough = obj._prepare_pr_answer() assert "AI summary" not in body assert f"### {PRDescriptionHeader.DIAGRAM_WALKTHROUGH.value}" in body @@ -367,7 +367,7 @@ def test_description_kept_by_default(self, mock_get_settings): mock_get_settings.return_value = _settings() obj = self._obj({"title": "t", "description": "AI summary"}) - _, body, _, _ = obj._prepare_pr_answer() + _, body, _ = obj._prepare_pr_answer() assert "AI summary" in body @@ -379,7 +379,7 @@ def test_description_list_value_is_joined_and_bullets_spaced(self, mock_get_sett "description": "Intro\n- one\n- two", }) - _, body, _, _ = obj._prepare_pr_answer() + _, body, _ = obj._prepare_pr_answer() # Bullet readability: single newline before "-" becomes double newline. assert "Intro\n\n- one\n\n- two" in body @@ -390,7 +390,7 @@ def test_diagram_section_uses_header_enum(self, mock_get_settings): diagram = "\n```mermaid\ngraph LR\nA --> B\n```" obj = self._obj({"title": "t", "description": "d", "changes_diagram": diagram}) - _, body, _, _ = obj._prepare_pr_answer() + _, body, _ = obj._prepare_pr_answer() assert f"### {PRDescriptionHeader.DIAGRAM_WALKTHROUGH.value}" in body assert "```mermaid" in body @@ -400,7 +400,7 @@ def test_title_uses_vars_title_when_data_has_no_title(self, mock_get_settings): mock_get_settings.return_value = _settings(generate_ai_title=False) obj = self._obj({"description": "d"}) - title, _, _, _ = obj._prepare_pr_answer() + title, _, _ = obj._prepare_pr_answer() assert title == "Original title" @@ -423,10 +423,9 @@ def test_non_gfm_provider_skips_table_rendering(self, mock_get_settings): obj = self._obj(gfm=False) value = {"backend": [("src/app.py", "Add cache", "Adds a bounded cache.")]} - body, comments = obj.process_pr_files_prediction("PRE", value) + body = obj.process_pr_files_prediction("PRE", value) assert body == "PRE" - assert comments == [] @patch("pr_agent.tools.pr_description.get_settings") def test_gfm_provider_emits_table_with_file_row(self, mock_get_settings): @@ -437,14 +436,13 @@ def test_gfm_provider_emits_table_with_file_row(self, mock_get_settings): obj = self._obj(gfm=True, diff_files=[diff]) value = {"backend": [("src/app.py", "Add cache", "Adds a bounded cache.")]} - body, comments = obj.process_pr_files_prediction("", value) + body = obj.process_pr_files_prediction("", value) assert body.startswith("
      • ") assert body.rstrip().endswith("
      ") assert "Backend" in body assert "app.py" in body assert "+5/-2" in body - assert comments == [] @patch("pr_agent.tools.pr_description.get_settings") def test_adaptive_collapsible_triggers_above_threshold(self, mock_get_settings): @@ -458,7 +456,7 @@ def test_adaptive_collapsible_triggers_above_threshold(self, mock_get_settings): ] } - body, _ = obj.process_pr_files_prediction("", value) + body = obj.process_pr_files_prediction("", value) assert "
      2 files" in body @@ -480,7 +478,7 @@ def test_walkthrough_table_round_trips_through_process_description(self, mock_ge obj.git_provider.get_line_link.return_value = "https://example/blob/main/src/app.py#L1" value = {"backend": [("src/app.py", "Add cache", "Adds a bounded cache.")]} - table, _ = obj.process_pr_files_prediction("", value) + table = obj.process_pr_files_prediction("", value) full_description = ( "Some intro text.\n\n___\n\n" diff --git a/tests/unittest/test_pr_help_docs_path_containment.py b/tests/unittest/test_pr_help_docs_path_containment.py new file mode 100644 index 0000000000..43007102c1 --- /dev/null +++ b/tests/unittest/test_pr_help_docs_path_containment.py @@ -0,0 +1,80 @@ +from types import SimpleNamespace + +from pr_agent.tools.pr_help_docs import PRHelpDocs + + +def _help_docs_with_clone(tmp_path, docs_path, include_root_readme_file=False): + clone_root = tmp_path / "clone" + clone_root.mkdir(parents=True, exist_ok=True) + (clone_root / "docs" / "guide.md").parent.mkdir(parents=True, exist_ok=True) + (clone_root / "docs" / "guide.md").write_text("docs content", encoding="utf-8") + provider = PRHelpDocs.__new__(PRHelpDocs) + provider.repo_url = "https://github.com/org/repo" + provider.ctx_url = "https://github.com/org/repo/pull/1" + provider.include_root_readme_file = include_root_readme_file + provider.supported_doc_exts = [".md"] + provider.docs_path = docs_path + provider.git_provider = SimpleNamespace(clone=lambda url, dst, remove_dest_folder: SimpleNamespace( + path=str(clone_root))) + return provider + + +def test_docs_path_within_clone_is_read(tmp_path): + provider = _help_docs_with_clone(tmp_path, "docs") + result = provider._gen_filenames_to_contents_map_from_repo() + assert result and "docs content" in next(iter(result.values())) + + +def test_absolute_docs_path_escaping_clone_is_rejected(tmp_path): + outside = tmp_path / "outside" + outside.mkdir() + (outside / "secret.md").write_text("HOST SECRET CONTENT", encoding="utf-8") + provider = _help_docs_with_clone(tmp_path, str(outside)) + assert provider._gen_filenames_to_contents_map_from_repo() == {} + + +def test_traversing_docs_path_escaping_clone_is_rejected(tmp_path): + outside = tmp_path / "outside" + outside.mkdir() + (outside / "secret.md").write_text("HOST SECRET CONTENT", encoding="utf-8") + provider = _help_docs_with_clone(tmp_path, "../outside") + assert provider._gen_filenames_to_contents_map_from_repo() == {} + + +def test_symlinked_doc_file_escaping_clone_is_skipped(tmp_path): + secret = tmp_path / "outside" / "secret.md" + secret.parent.mkdir(parents=True, exist_ok=True) + secret.write_text("HOST SECRET CONTENT", encoding="utf-8") + provider = _help_docs_with_clone(tmp_path, "docs") + (tmp_path / "clone" / "docs" / "leak.md").symlink_to(secret) + + result = provider._gen_filenames_to_contents_map_from_repo() + + assert any("docs content" in v for v in result.values()) + assert not any("HOST SECRET" in v for v in result.values()) + + +def test_symlinked_root_readme_escaping_clone_is_skipped(tmp_path): + secret = tmp_path / "outside" / "secret.md" + secret.parent.mkdir(parents=True, exist_ok=True) + secret.write_text("HOST SECRET CONTENT", encoding="utf-8") + provider = _help_docs_with_clone(tmp_path, "docs", include_root_readme_file=True) + clone_root = tmp_path / "clone" + (clone_root / "readme.txt").write_text("clone readme", encoding="utf-8") + (clone_root / "README.md").symlink_to(secret) + + result = provider._gen_filenames_to_contents_map_from_repo() + + assert any("clone readme" in v for v in result.values()) + assert not any("HOST SECRET" in v for v in result.values()) + + +def test_symlink_to_in_clone_file_is_kept(tmp_path): + provider = _help_docs_with_clone(tmp_path, "docs") + clone_root = tmp_path / "clone" + (clone_root / "docs" / "guide.md").write_text("original content", encoding="utf-8") + (clone_root / "docs" / "alias.md").symlink_to(clone_root / "docs" / "guide.md") + + result = provider._gen_filenames_to_contents_map_from_repo() + + assert sum("original content" in v for v in result.values()) >= 1 diff --git a/tests/unittest/test_pr_help_docs_row_isolation.py b/tests/unittest/test_pr_help_docs_row_isolation.py new file mode 100644 index 0000000000..4073382a76 --- /dev/null +++ b/tests/unittest/test_pr_help_docs_row_isolation.py @@ -0,0 +1,173 @@ +"""Regression tests for malformed model-produced /help_docs collections.""" + +import pytest + +from pr_agent.tools.pr_help_docs import ( + PRHelpDocs, + format_markdown_q_and_a_response, + get_valid_ranking_indices, + modify_answer_section, +) + + +class StubProvider: + def get_canonical_url_parts(self, **kwargs): + return "https://example.com/repo/blob/main", "" + + +def test_format_markdown_q_and_a_response_preserves_valid_sources_when_one_row_is_malformed(): + answer = format_markdown_q_and_a_response( + "Where is the guide?", + "The answer is here.", + [ + {"file_name": "/docs/guide.md", "relevant_section_header_string": "Getting started"}, + {"file_name": "/docs/broken.md"}, + ], + [".md"], + "https://example.com/repo/blob/main", + ) + + assert "The answer is here." in answer + assert "https://example.com/repo/blob/main/docs/guide.md#getting-started" in answer + assert "broken.md" not in answer + + +def test_format_markdown_q_and_a_response_omits_sources_when_all_rows_are_malformed(): + answer = format_markdown_q_and_a_response( + "Where is the guide?", + "The answer is here.", + [{"file_name": "/docs/broken.md"}, None, {"relevant_section_header_string": "Missing file"}], + [".md"], + "https://example.com/repo/blob/main", + ) + + assert answer == "### Question: \nWhere is the guide?\n\n### Answer:\nThe answer is here.\n\n" + assert "Relevant Sources" not in answer + + +def test_format_markdown_q_and_a_response_omits_malformed_source_collection(): + answer = format_markdown_q_and_a_response( + "Where is the guide?", + "The answer is here.", + {"file_name": "/docs/guide.md", "relevant_section_header_string": "Getting started"}, + [".md"], + "https://example.com/repo/blob/main", + ) + + assert answer == "### Question: \nWhere is the guide?\n\n### Answer:\nThe answer is here.\n\n" + + +def test_format_markdown_q_and_a_response_routes_unexpected_formatter_errors_to_failure_handler(monkeypatch): + def fail_to_format_header(header): + raise RuntimeError("unexpected formatter failure") + + monkeypatch.setattr("pr_agent.tools.pr_help_docs.format_markdown_header", fail_to_format_header) + + answer = format_markdown_q_and_a_response( + "Where is the guide?", + "The answer is here.", + [{"file_name": "/docs/guide.md", "relevant_section_header_string": "Getting started"}], + [".md"], + "https://example.com/repo/blob/main", + ) + + assert answer == "" + + +def test_modify_answer_section_preserves_answer_without_sources(): + assert modify_answer_section("### Answer:\nThe answer is here.\n\n") == ( + "### :bulb: Auto-generated documentation-based answer:\nThe answer is here.\n\n" + ) + + +def test_format_model_answer_preserves_answer_when_no_sources_survive(): + tool = PRHelpDocs.__new__(PRHelpDocs) + tool.question = "Where is the guide?" + tool.supported_doc_exts = [".md"] + tool.return_as_string = True + tool.repo_url = "https://example.com/org/repo" + tool.repo_url_given_explicitly = True + tool.repo_desired_branch = "main" + tool.git_provider = StubProvider() + + assert tool._format_model_answer("The answer is here.", [{"file_name": "/docs/broken.md"}]) == ( + "### :bulb: Auto-generated documentation-based answer:\nThe answer is here.\n\n" + ) + + +@pytest.mark.asyncio +async def test_run_preserves_answer_when_model_returns_empty_source_list(monkeypatch): + async def fake_retry_with_fallback_models(*args, **kwargs): + return """question_is_relevant: 1\nresponse: The answer is here.\nrelevant_sections: []\n""" + + monkeypatch.setattr("pr_agent.tools.pr_help_docs.retry_with_fallback_models", fake_retry_with_fallback_models) + + tool = PRHelpDocs.__new__(PRHelpDocs) + tool.question = "Where is the guide?" + tool.return_as_string = True + tool.repo_url = "https://example.com/org/repo" + tool.repo_url_given_explicitly = True + tool.repo_desired_branch = "main" + tool.supported_doc_exts = [".md"] + tool.git_provider = StubProvider() + tool.ai_handler = object() + tool.vars = {"question": tool.question, "snippets": ""} + tool._gen_filenames_to_contents_map_from_repo = lambda: {"/docs/guide.md": "# Guide\n\nGuide content."} + tool._trim_docs_input = lambda docs_input, *args, only_return_if_trim_needed=False, **kwargs: ( + False if only_return_if_trim_needed else docs_input + ) + + assert await tool.run() == ( + "### :bulb: Auto-generated documentation-based answer:\nThe answer is here.\n\n" + ) + + +def test_get_valid_ranking_indices_preserves_valid_order_and_skips_malformed_rows(): + assert get_valid_ranking_indices( + [ + {"idx": "2"}, + {"idx": True}, + {"idx": "not-a-number"}, + {"idx": 0}, + {"idx": 2}, + {"idx": 4}, + {"idx": 1.0}, + "not-a-row", + ], + 3, + ) == [2, 0, 2] + + +def test_get_valid_ranking_indices_skips_non_list_collections(): + assert get_valid_ranking_indices({"idx": 0}, 1) == [] + assert get_valid_ranking_indices(True, 1) == [] + + +def test_get_valid_ranking_indices_skips_all_malformed_rows(): + assert get_valid_ranking_indices( + [{"idx": "not-a-number"}, {"idx": -1}, {"idx": 2}, None], + 2, + ) == [] + + +@pytest.mark.asyncio +async def test_rank_docs_preserves_valid_rows_when_model_ranking_is_partially_malformed(monkeypatch): + async def fake_retry_with_fallback_models(*args, **kwargs): + return """relevant_files_ranking:\n - idx: 2\n - invalid: row\n - idx: 0\n""" + + monkeypatch.setattr("pr_agent.tools.pr_help_docs.retry_with_fallback_models", fake_retry_with_fallback_models) + + tool = PRHelpDocs.__new__(PRHelpDocs) + tool.ai_handler = object() + tool.vars = {"question": "Where is the guide?", "snippets": ""} + tool._trim_docs_input = lambda docs_input, *args, **kwargs: docs_input + docs = { + "/docs/first.md": "# First\n\nFirst document.", + "/docs/second.md": "# Second\n\nSecond document.", + "/docs/third.md": "# Third\n\nThird document.", + } + + ranked_prompt = await tool._rank_docs_and_return_them_as_prompt(docs, 10_000) + + assert "/docs/second.md" not in ranked_prompt + assert ranked_prompt.index("/docs/third.md") < ranked_prompt.index("/docs/first.md") diff --git a/tests/unittest/test_pr_help_message_fallback.py b/tests/unittest/test_pr_help_message_fallback.py new file mode 100644 index 0000000000..c6bb6ec1ed --- /dev/null +++ b/tests/unittest/test_pr_help_message_fallback.py @@ -0,0 +1,471 @@ +"""Question-mode help delegates failed attempts to the shared fallback loop.""" + +import asyncio +from math import ceil +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from pr_agent.algo import run_details +from pr_agent.config_loader import get_settings +from pr_agent.tools import pr_help_message +from pr_agent.tools.pr_help_message import PRHelpMessage +from tests.unittest._settings_helpers import restore_settings, snapshot_settings + +PRIMARY = "gpt-4o" +BACKUP = "gpt-4o-mini" +ANSWER = ( + "response: Enable automatic review in the repository settings.\n" + "relevant_sections:\n" + " - file_name: /tools/review.md\n" + " relevant_section_header_string: Automatic review\n" +) + + +@pytest.fixture +def help_tool(tmp_path, monkeypatch): + overrides = { + "config.model": PRIMARY, + "config.fallback_models": [BACKUP], + "config.publish_output": True, + "config.max_output_tokens": 0, + "config.model_token_count_estimate_factor": 0, + "model_routing.enable": False, + # Keep this regression independent of the separate OpenAI-key gate fix. + "openai.key": "test-only-key", + "openai.deployment_id": "primary-deployment", + "openai.fallback_deployments": ["backup-deployment"], + } + snapshot = snapshot_settings([*overrides, "pr_help_prompts.system", "pr_help_prompts.user"]) + for key, value in overrides.items(): + get_settings().set(key, value) + + doc = tmp_path / "docs" / "docs" / "tools" / "review.md" + doc.parent.mkdir(parents=True) + doc.write_text("# Automatic review\n\nEnable automatic review in the repository settings.", encoding="utf-8") + monkeypatch.setattr(pr_help_message, "__file__", str(tmp_path / "pr_agent" / "tools" / "pr_help_message.py")) + + details = run_details.RunDetails() + monkeypatch.setattr(run_details, "get_run_details", lambda: details) + logger = Mock() + monkeypatch.setattr(pr_help_message, "get_logger", lambda: logger) + tool = PRHelpMessage.__new__(PRHelpMessage) + tool.git_provider = SimpleNamespace(pr_url="https://example.com/org/repo/pull/1", publish_comment=Mock()) + tool.ai_handler = SimpleNamespace(chat_completion=AsyncMock()) + tool.question_str = "How do I configure automatic reviews?" + tool.return_as_string = False + tool.vars = {"question": tool.question_str, "snippets": ""} + tool._fixture_doc = doc + try: + yield tool, details, logger + finally: + restore_settings(snapshot) + + +def attempted_models(tool): + return [call.kwargs["model"] for call in tool.ai_handler.chat_completion.await_args_list] + + +def count_message_characters(*, messages, **_kwargs): + return sum(len(message["content"]) for message in messages) + + +async def test_model_prompt_preserves_question_and_documentation_markup(help_tool): + tool, _, _ = help_tool + question = "What does ` & 'notes'` mean?" + snippets = "Use `a < b && c > d` and ``." + tool.vars.update(question=question, snippets=snippets) + tool.ai_handler.chat_completion.return_value = ANSWER, "stop" + + assert await tool._prepare_prediction(PRIMARY) == ANSWER + + tool.ai_handler.chat_completion.assert_awaited_once() + user_prompt = tool.ai_handler.chat_completion.await_args.kwargs["user"] + assert question in user_prompt + assert snippets in user_prompt + + +async def test_primary_failure_uses_backup_answer(help_tool): + tool, details, _ = help_tool + tool.ai_handler.chat_completion.side_effect = [RuntimeError("primary unavailable"), (ANSWER, "stop")] + + await tool.run() + + assert attempted_models(tool) == [PRIMARY, BACKUP] + backup_prompt = tool.ai_handler.chat_completion.await_args_list[1].kwargs["user"] + assert tool.question_str in backup_prompt + assert "Enable automatic review in the repository settings." in backup_prompt + tool.git_provider.publish_comment.assert_called_once() + published_comment = tool.git_provider.publish_comment.call_args.args[0] + assert "### Answer:\nEnable automatic review" in published_comment + assert "> - https://docs.pr-agent.ai/tools/review/#automatic-review" in published_comment + assert details.model_used == BACKUP + assert details.fallback_used is True + assert get_settings().get("openai.deployment_id") == "primary-deployment" + + +async def test_each_attempt_fits_complete_prompt_for_its_model(help_tool, monkeypatch): + tool, details, _ = help_tool + tool._fixture_doc.write_text("# Automatic review\n\n" + "A" * 600, encoding="utf-8") + get_settings().set("pr_help_prompts.system", "system") + get_settings().set("pr_help_prompts.user", "question={{ question }}\nsnippets={{ snippets }}") + get_settings().set("config.max_output_tokens", 20) + limits = {PRIMARY: 2_000, BACKUP: 260} + limit_calls = [] + counted_models = [] + + def model_limit(model, ignore_max_model_tokens=False): + limit_calls.append((model, ignore_max_model_tokens)) + return limits[model] + + def model_token_count(*, model, messages): + counted_models.append(model) + character_count = count_message_characters(messages=messages) + return character_count if model == PRIMARY else character_count * 2 + + monkeypatch.setattr(pr_help_message, "get_max_tokens", model_limit) + monkeypatch.setattr(pr_help_message, "token_counter", model_token_count) + tool.ai_handler.chat_completion.side_effect = [RuntimeError("primary unavailable"), (ANSWER, "stop")] + + await tool.run() + + assert attempted_models(tool) == [PRIMARY, BACKUP] + primary_prompt = tool.ai_handler.chat_completion.await_args_list[0].kwargs["user"] + backup_prompt = tool.ai_handler.chat_completion.await_args_list[1].kwargs["user"] + assert "A" * 600 in primary_prompt + assert "...(truncated)" not in primary_prompt + assert "A" * 600 not in backup_prompt + assert "...(truncated)" in backup_prompt + assert "==file name==" in backup_prompt + assert "/tools/review.md" in backup_prompt + assert 2 * (len(backup_prompt) + len("system")) <= limits[BACKUP] - 20 + assert "A" * 600 in tool.vars["snippets"] + assert limit_calls == [(PRIMARY, True), (BACKUP, True)] + assert PRIMARY in counted_models + assert BACKUP in counted_models + assert details.model_used == BACKUP + + +@pytest.mark.parametrize( + ("configured_output", "expected_budget"), + [ + (400, 2_600), + ("400", 2_600), + (0, 1_000), + (-1, 1_000), + (None, 1_000), + ("invalid", 1_000), + (400.0, 2_600), + (400.5, 2_600), + (float("nan"), 1_000), + (float("inf"), 1_000), + (True, 2_999), + (False, 1_000), + ("1" * 5_000, 1_000), + ], +) +def test_prompt_budget_honors_positive_output_limit_or_help_default( + help_tool, monkeypatch, configured_output, expected_budget +): + tool, _, _ = help_tool + get_settings().set("config.max_output_tokens", configured_output) + monkeypatch.setattr(pr_help_message, "get_max_tokens", lambda *_args, **_kwargs: 3_000) + + assert tool._get_prompt_budget(PRIMARY) == expected_budget + + +def test_prompt_budget_uses_handler_reported_output_limit(help_tool, monkeypatch): + tool, _, _ = help_tool + get_settings().set("config.max_output_tokens", 400) + tool.ai_handler.get_output_token_limit = Mock(return_value=1_600) + monkeypatch.setattr(pr_help_message, "get_max_tokens", lambda *_args, **_kwargs: 3_000) + + assert tool._get_prompt_budget(PRIMARY) == 1_400 + tool.ai_handler.get_output_token_limit.assert_called_once_with(PRIMARY) + + +def test_prompt_budget_uses_handler_reported_output_reserve(help_tool, monkeypatch): + tool, _, _ = help_tool + get_settings().set("config.max_output_tokens", 400) + tool.ai_handler.get_output_token_reserve = Mock(return_value=1_600) + monkeypatch.setattr(pr_help_message, "get_max_tokens", lambda *_args, **_kwargs: 3_000) + + assert tool._get_prompt_budget(PRIMARY) == 1_400 + tool.ai_handler.get_output_token_reserve.assert_called_once_with(PRIMARY, 2_000) + + +@pytest.mark.parametrize("reported_limit", [None, 0, -1, True, 1.5, "1600"]) +def test_unusable_handler_output_limit_falls_back_to_config( + help_tool, monkeypatch, reported_limit +): + tool, _, _ = help_tool + get_settings().set("config.max_output_tokens", 400) + tool.ai_handler.get_output_token_limit = Mock(return_value=reported_limit) + monkeypatch.setattr(pr_help_message, "get_max_tokens", lambda *_args, **_kwargs: 3_000) + + assert tool._get_prompt_budget(PRIMARY) == 2_600 + + +def test_failing_handler_output_limit_falls_back_to_config(help_tool, monkeypatch): + tool, _, logger = help_tool + get_settings().set("config.max_output_tokens", 400) + tool.ai_handler.get_output_token_limit = Mock(side_effect=RuntimeError("unavailable")) + monkeypatch.setattr(pr_help_message, "get_max_tokens", lambda *_args, **_kwargs: 3_000) + + assert tool._get_prompt_budget(PRIMARY) == 2_600 + assert any("output token limit" in call.args[0] for call in logger.debug.call_args_list) + + +@pytest.mark.parametrize( + ("estimate_factor", "expected"), + [ + (-0.5, 53), + (0, 53), + (None, 53), + ("invalid", 53), + (float("nan"), 53), + (float("inf"), 53), + (True, 53), + (0.3, ceil(53 * 1.3)), + ], +) +def test_local_prompt_estimate_never_reduces_encoded_content( + help_tool, monkeypatch, estimate_factor, expected +): + tool, _, _ = help_tool + + class CharacterEncoder: + @staticmethod + def encode(text, disallowed_special=()): + return list(text) + + get_settings().set("config.model_token_count_estimate_factor", estimate_factor) + monkeypatch.setattr(pr_help_message, "token_counter", Mock(side_effect=RuntimeError("counter unavailable")), + raising=False) + monkeypatch.setattr(pr_help_message.TokenEncoder, "get_token_encoder", lambda _model: CharacterEncoder(), + raising=False) + + assert tool._count_prompt_tokens(PRIMARY, "abc", "de") == expected + + +def test_local_prompt_estimate_uses_raw_estimate_on_overflow(help_tool, monkeypatch): + tool, _, _ = help_tool + + class CharacterEncoder: + @staticmethod + def encode(text, disallowed_special=()): + return list(text) + + get_settings().set("config.model_token_count_estimate_factor", 1e308) + monkeypatch.setattr(pr_help_message, "token_counter", Mock(side_effect=RuntimeError("counter unavailable"))) + monkeypatch.setattr(pr_help_message.TokenEncoder, "get_token_encoder", lambda _model: CharacterEncoder()) + + assert tool._count_prompt_tokens(PRIMARY, "abc", "de") == 53 + + +@pytest.mark.parametrize("counter_mode", ["error", "zero", "boolean", "non_integer"]) +async def test_unusable_model_count_uses_local_estimate_before_provider_call( + help_tool, monkeypatch, counter_mode +): + tool, _, _ = help_tool + tool.vars = {"question": "q", "snippets": "A" * 100} + get_settings().set("pr_help_prompts.system", "") + get_settings().set("pr_help_prompts.user", "{{ snippets }}") + monkeypatch.setattr(tool, "_get_prompt_budget", lambda _model: 70) + + class CharacterEncoder: + @staticmethod + def encode(text, disallowed_special=()): + return list(text) + + def unusable_count(**_kwargs): + if counter_mode == "error": + raise RuntimeError("counter unavailable") + return {"zero": 0, "boolean": True, "non_integer": 1.5}[counter_mode] + + monkeypatch.setattr(pr_help_message, "token_counter", unusable_count) + monkeypatch.setattr(pr_help_message.TokenEncoder, "get_token_encoder", lambda _model: CharacterEncoder()) + tool.ai_handler.chat_completion.return_value = ANSWER, "stop" + + await tool._prepare_prediction(PRIMARY) + + tool.ai_handler.chat_completion.assert_awaited_once() + user_prompt = tool.ai_handler.chat_completion.await_args.kwargs["user"] + assert "...(truncated)" in user_prompt + assert len(user_prompt) + 48 <= 70 + + +async def test_local_estimate_overflow_still_tries_providers(help_tool, monkeypatch): + tool, details, _ = help_tool + + class CharacterEncoder: + @staticmethod + def encode(text, disallowed_special=()): + return list(text) + + get_settings().set("config.model_token_count_estimate_factor", 1e308) + monkeypatch.setattr(pr_help_message, "token_counter", Mock(side_effect=RuntimeError("counter unavailable"))) + monkeypatch.setattr(pr_help_message.TokenEncoder, "get_token_encoder", lambda _model: CharacterEncoder()) + tool._prepare_prediction = AsyncMock(wraps=tool._prepare_prediction) + tool.ai_handler.chat_completion.side_effect = [RuntimeError("primary unavailable"), (ANSWER, "stop")] + + await tool.run() + + assert [call.args[0] for call in tool._prepare_prediction.await_args_list] == [PRIMARY, BACKUP] + assert attempted_models(tool) == [PRIMARY, BACKUP] + tool.git_provider.publish_comment.assert_called_once() + assert details.model_used == BACKUP + + +async def test_empty_claude_system_prompt_is_normalized_before_fitting(help_tool, monkeypatch): + tool, _, _ = help_tool + model = "anthropic/claude-sonnet-4-5" + tool.vars = {"question": "q", "snippets": "A" * 100} + get_settings().set("pr_help_prompts.system", "") + get_settings().set("pr_help_prompts.user", "{{ snippets }}") + monkeypatch.setattr(tool, "_get_prompt_budget", lambda _model: 80) + tool.ai_handler.normalize_request_prompts = pr_help_message.LiteLLMAIHandler.normalize_request_prompts + counted_system_prompts = [] + + def count_normalized_prompts(*, messages, **_kwargs): + counted_system_prompts.append(messages[0]["content"]) + return count_message_characters(messages=messages) + + monkeypatch.setattr(pr_help_message, "token_counter", count_normalized_prompts) + tool.ai_handler.chat_completion.return_value = ANSWER, "stop" + + await tool._prepare_prediction(model) + + assert counted_system_prompts + assert set(counted_system_prompts) == {"No system prompt provided"} + call = tool.ai_handler.chat_completion.await_args + assert call.kwargs["system"] == "No system prompt provided" + assert "...(truncated)" in call.kwargs["user"] + + +def test_prompt_fitting_checks_smallest_nonempty_prefix(help_tool, monkeypatch): + tool, _, _ = help_tool + get_settings().set("pr_help_prompts.system", "") + get_settings().set("pr_help_prompts.user", "{{ snippets }}") + monkeypatch.setattr(tool, "_get_prompt_budget", lambda _model: 10) + + def discontinuous_count(*, messages, **_kwargs): + snippets = messages[1]["content"] + if snippets in ("", pr_help_message.TRUNCATION_MARKER): + return 1 + if snippets == f"A{pr_help_message.TRUNCATION_MARKER}": + return 10 + return 1_000 + + monkeypatch.setattr(pr_help_message, "token_counter", discontinuous_count) + + _, user_prompt = tool._fit_prompts({"question": "q", "snippets": "ABC"}, PRIMARY) + + assert user_prompt == f"A{pr_help_message.TRUNCATION_MARKER}" + + +async def test_marker_only_prompt_is_sent_when_no_document_character_fits(help_tool, monkeypatch): + tool, _, _ = help_tool + tool.vars = {"question": "q", "snippets": "ABC"} + get_settings().set("pr_help_prompts.system", "") + get_settings().set("pr_help_prompts.user", "{{ snippets }}") + monkeypatch.setattr(tool, "_get_prompt_budget", lambda _model: 10) + + def marker_only_count(*, messages, **_kwargs): + snippets = messages[1]["content"] + if snippets == "": + return 1 + if snippets == pr_help_message.TRUNCATION_MARKER: + return 5 + return 1_000 + + monkeypatch.setattr(pr_help_message, "token_counter", marker_only_count) + tool.ai_handler.chat_completion.return_value = ANSWER, "stop" + + await tool._prepare_prediction(PRIMARY) + + tool.ai_handler.chat_completion.assert_awaited_once() + assert tool.ai_handler.chat_completion.await_args.kwargs["user"] == pr_help_message.TRUNCATION_MARKER + + +async def test_fixed_prompt_overhead_skips_model_call_and_tries_larger_fallback(help_tool, monkeypatch): + tool, details, _ = help_tool + get_settings().set("pr_help_prompts.system", "S" * 300) + get_settings().set("pr_help_prompts.user", "question={{ question }}\nsnippets={{ snippets }}") + get_settings().set("config.max_output_tokens", 20) + limits = {PRIMARY: 250, BACKUP: 2_000} + + monkeypatch.setattr(pr_help_message, "get_max_tokens", + lambda model, ignore_max_model_tokens=False: limits[model]) + monkeypatch.setattr(pr_help_message, "token_counter", count_message_characters, raising=False) + tool._prepare_prediction = AsyncMock(wraps=tool._prepare_prediction) + tool.ai_handler.chat_completion.return_value = ANSWER, "stop" + + await tool.run() + + assert [call.args[0] for call in tool._prepare_prediction.await_args_list] == [PRIMARY, BACKUP] + assert attempted_models(tool) == [BACKUP] + assert details.model_used == BACKUP + + +async def test_all_models_fail_without_publishing_no_information(help_tool): + tool, details, logger = help_tool + tool.ai_handler.chat_completion.side_effect = RuntimeError("provider unavailable") + + assert await tool.run() == "" + + assert attempted_models(tool) == [PRIMARY, BACKUP] + tool.git_provider.publish_comment.assert_not_called() + logger.exception.assert_called_once() + assert "Failed to generate prediction with any model" in logger.exception.call_args.args[0] + assert details.model_used is None + assert get_settings().get("openai.deployment_id") == "primary-deployment" + + +@pytest.mark.parametrize( + "response, expected", + [(ANSWER, "### Answer:\nEnable automatic review"), + ("response: No matching information\nrelevant_sections: []", "Could not find relevant information")], +) +async def test_successful_primary_response_does_not_use_backup(help_tool, response, expected): + tool, details, _ = help_tool + tool.ai_handler.chat_completion.return_value = response, "stop" + + await tool.run() + + assert attempted_models(tool) == [PRIMARY] + tool.git_provider.publish_comment.assert_called_once() + assert expected in tool.git_provider.publish_comment.call_args.args[0] + assert details.model_used == PRIMARY + assert details.fallback_used is False + + +async def test_cancellation_propagates_without_backup(help_tool): + tool, details, logger = help_tool + tool.ai_handler.chat_completion.side_effect = asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await tool.run() + + assert attempted_models(tool) == [PRIMARY] + tool.git_provider.publish_comment.assert_not_called() + logger.exception.assert_not_called() + assert details.model_used is None + assert get_settings().get("openai.deployment_id") == "primary-deployment" + + +async def test_invalid_prompt_reaches_final_failure_boundary(help_tool): + tool, details, logger = help_tool + get_settings().set("pr_help_prompts.user", "{{ missing_variable }}") + tool._prepare_prediction = AsyncMock(wraps=tool._prepare_prediction) + + assert await tool.run() == "" + + assert [call.args[0] for call in tool._prepare_prediction.await_args_list] == [PRIMARY, BACKUP] + tool.ai_handler.chat_completion.assert_not_called() + tool.git_provider.publish_comment.assert_not_called() + logger.exception.assert_called_once() + assert "Failed to generate prediction with any model" in logger.exception.call_args.args[0] + assert details.model_used is None + assert get_settings().get("openai.deployment_id") == "primary-deployment" diff --git a/tests/unittest/test_pr_help_message_rendering.py b/tests/unittest/test_pr_help_message_rendering.py new file mode 100644 index 0000000000..d00d1e2fd5 --- /dev/null +++ b/tests/unittest/test_pr_help_message_rendering.py @@ -0,0 +1,222 @@ +"""Regression tests for provider-independent /help behavior.""" + +import pytest + +from pr_agent.config_loader import get_settings +from pr_agent.git_providers.bitbucket_provider import BitbucketProvider +from pr_agent.git_providers.bitbucket_server_provider import BitbucketServerProvider +from pr_agent.git_providers.github_provider import GithubProvider +from pr_agent.git_providers.gitlab_provider import GitLabProvider +from pr_agent.tools import pr_help_message as pr_help_message_module +from pr_agent.tools.pr_help_message import PRHelpMessage +from tests.unittest._settings_helpers import restore_settings, snapshot_settings + +INTERACTIVE_MARKER = "Trigger Interactively" +PLAIN_TABLE_MARKER = "| Tool | Description |" +UNSUPPORTED_MARKER = "requires gfm markdown" +CURRENT_DOCS_URL = "https://docs.pr-agent.ai" +RETIRED_DOCS_HOSTS = ("qodo-merge-docs.qodo.ai", "pr-agent-docs.codium.ai") + + +class StubProvider: + """A provider that is none of the concrete classes the tool used to branch on.""" + + def __init__(self, gfm_markdown: bool, markdown_tables: bool = False, checkbox_commands: bool = False): + self.pr_url = "https://example.com/org/repo/pull/1" + self._gfm_markdown = gfm_markdown + self._markdown_tables = markdown_tables + self._checkbox_commands = checkbox_commands + self.published = [] + + def is_supported(self, capability: str) -> bool: + return self._gfm_markdown if capability == "gfm_markdown" else True + + def supports_markdown_tables(self) -> bool: + return self._markdown_tables + + def supports_checkbox_commands(self) -> bool: + return self._checkbox_commands + + def publish_comment(self, pr_comment: str, is_temporary: bool = False): + self.published.append(pr_comment) + + +class StubAiHandler: + def __init__(self, response=None, error=None): + self.response = response + self.error = error + self.calls = [] + + async def chat_completion(self, *, model, temperature, system, user): + self.calls.append({"model": model, "system": system, "user": user}) + if self.error: + raise self.error + return self.response, "stop" + + +class StubTokenHandler: + @staticmethod + def count_tokens(text): + return len(text) + + +async def run_walkthrough(provider) -> str: + tool = PRHelpMessage.__new__(PRHelpMessage) + tool.git_provider = provider + tool.question_str = "" + tool.return_as_string = False + await tool.run() + assert len(provider.published) == 1 + return provider.published[0] + + +def build_question_tool(tmp_path, monkeypatch, handler): + package_root = tmp_path / "package" + docs_path = package_root / "docs" / "docs" + review_doc = docs_path / "tools" / "review.md" + review_doc.parent.mkdir(parents=True) + review_doc.write_text("# Automatic review\n\nEnable automatic review in the repository settings.", encoding="utf-8") + + source_path = package_root / "pr_agent" / "tools" / "pr_help_message.py" + monkeypatch.setattr(pr_help_message_module, "Path", lambda _: source_path) + + tool = PRHelpMessage.__new__(PRHelpMessage) + tool.git_provider = StubProvider(gfm_markdown=True) + tool.ai_handler = handler + tool.question_str = "How do I configure automatic reviews?" + tool.return_as_string = False + tool.vars = {"question": tool.question_str, "snippets": ""} + tool.token_handler = StubTokenHandler() + return tool + + +@pytest.fixture +def published_output(): + snapshot = snapshot_settings(["config.publish_output", "config.disable_checkboxes"]) + get_settings().set("config.publish_output", True) + yield + restore_settings(snapshot) + + +@pytest.fixture +def non_openai_question_settings(): + keys = [ + "config.model", + "config.fallback_models", + "model_routing.enable", + "openai.key", + "openai.deployment_id", + ] + snapshot = snapshot_settings(keys) + settings = get_settings() + settings.set("config.model", "anthropic/claude-3-5-sonnet-20240620") + settings.set("config.fallback_models", []) + settings.set("model_routing.enable", False) + settings.set("openai.key", None) + settings.set("openai.deployment_id", None) + yield + restore_settings(snapshot) + + +async def test_question_reaches_configured_handler_without_openai_key( + published_output, non_openai_question_settings, tmp_path, monkeypatch +): + handler = StubAiHandler( + response=( + "response: Enable automatic review in the repository settings.\n" + "relevant_sections:\n" + " - file_name: /tools/review.md\n" + " relevant_section_header_string: Automatic review\n" + ) + ) + tool = build_question_tool(tmp_path, monkeypatch, handler) + + await tool.run() + + assert [call["model"] for call in handler.calls] == ["anthropic/claude-3-5-sonnet-20240620"] + assert tool.question_str in handler.calls[0]["user"] + assert "Enable automatic review in the repository settings." in handler.calls[0]["user"] + assert "Enable automatic review in the repository settings." in tool.git_provider.published[0] + assert "requires an OpenAI API key" not in tool.git_provider.published[0] + + +async def test_question_uses_configured_handler_error_path_without_openai_key( + published_output, non_openai_question_settings, tmp_path, monkeypatch +): + handler = StubAiHandler(error=RuntimeError("provider credentials missing")) + tool = build_question_tool(tmp_path, monkeypatch, handler) + + await tool.run() + + assert [call["model"] for call in handler.calls] == ["anthropic/claude-3-5-sonnet-20240620"] + assert tool.git_provider.published == [] + + +@pytest.mark.parametrize( + "capabilities, expected, unexpected", + [ + # gfm markdown plus checkbox commands is the only combination that offers the checkboxes + (dict(gfm_markdown=True, checkbox_commands=True), INTERACTIVE_MARKER, PLAIN_TABLE_MARKER), + (dict(gfm_markdown=True), "", INTERACTIVE_MARKER), + # without gfm markdown, a provider that renders pipe tables gets the basic table + (dict(gfm_markdown=False, markdown_tables=True), PLAIN_TABLE_MARKER, "
      "), + (dict(gfm_markdown=False), UNSUPPORTED_MARKER, "
      "), + ], +) +async def test_rendering_follows_capabilities(published_output, capabilities, expected, unexpected): + comment = await run_walkthrough(StubProvider(**capabilities)) + assert expected in comment + assert unexpected not in comment + + +async def test_checkboxes_stay_disabled_by_configuration(published_output): + get_settings().set("config.disable_checkboxes", True) + comment = await run_walkthrough(StubProvider(gfm_markdown=True, checkbox_commands=True)) + assert INTERACTIVE_MARKER not in comment + assert "
      " in comment + + +@pytest.mark.parametrize("checkbox_commands", [False, True], ids=["without-checkboxes", "with-checkboxes"]) +async def test_walkthrough_uses_current_documentation_site(published_output, checkbox_commands): + comment = await run_walkthrough(StubProvider(gfm_markdown=True, checkbox_commands=checkbox_commands)) + + assert (INTERACTIVE_MARKER in comment) is checkbox_commands + assert f"{CURRENT_DOCS_URL}/tools/review/" in comment + assert f"{CURRENT_DOCS_URL}/usage-guide/automations_and_usage/" in comment + assert not any(host in comment for host in RETIRED_DOCS_HOSTS) + + +@pytest.mark.parametrize( + ("file_name", "header", "expected"), + [ + ("/index.md", "", f"{CURRENT_DOCS_URL}/"), + ("/tools/review.md", "Automatic review", f"{CURRENT_DOCS_URL}/tools/review/#automatic-review"), + ("/faq/index.md", "Frequently asked questions", f"{CURRENT_DOCS_URL}/faq/#frequently-asked-questions"), + ], +) +def test_question_source_urls_use_canonical_documentation_paths(file_name, header, expected): + tool = PRHelpMessage.__new__(PRHelpMessage) + + assert tool.format_docs_url(file_name, header) == expected + + +def test_question_system_prompt_identifies_the_open_source_pr_agent_project(): + system_prompt, _ = PRHelpMessage._render_prompts({"question": "What is PR-Agent?", "snippets": ""}) + + assert "the open-source PR-Agent project" in system_prompt + assert "recently renamed" not in system_prompt.lower() + + +@pytest.mark.parametrize( + "provider_class, checkbox_commands, markdown_tables", + [ + (GithubProvider, True, False), + (GitLabProvider, False, False), + (BitbucketProvider, False, False), + (BitbucketServerProvider, False, True), + ], +) +def test_providers_declare_the_capabilities_the_tool_reads(provider_class, checkbox_commands, markdown_tables): + provider = provider_class.__new__(provider_class) + assert provider.supports_checkbox_commands() is checkbox_commands + assert provider.supports_markdown_tables() is markdown_tables diff --git a/tests/unittest/test_pr_processing_core.py b/tests/unittest/test_pr_processing_core.py index 13ad4a9928..0c29243d19 100644 --- a/tests/unittest/test_pr_processing_core.py +++ b/tests/unittest/test_pr_processing_core.py @@ -13,22 +13,156 @@ class FakeTokenHandler: def __init__(self, prompt_tokens=100): self.prompt_tokens = prompt_tokens + self.count_calls = 0 def count_tokens(self, patch): + self.count_calls += 1 return len(patch.split()) class FakeProvider: def __init__(self, files): self.files = files + self.diff_calls = 0 + self.language_calls = 0 def get_diff_files(self): + self.diff_calls += 1 return self.files def get_languages(self): + self.language_calls += 1 return {"Python": 100} +def test_prepared_pr_diff_reuses_compressed_files_without_changing_chunks(monkeypatch): + settings = get_settings() + original = { + "patch_extra_lines_before": settings.config.patch_extra_lines_before, + "patch_extra_lines_after": settings.config.patch_extra_lines_after, + "large_patch_policy": settings.config.get("large_patch_policy", "skip"), + "verbosity_level": settings.config.verbosity_level, + } + settings.config.patch_extra_lines_before = 0 + settings.config.patch_extra_lines_after = 0 + settings.config.large_patch_policy = "skip" + settings.config.verbosity_level = 0 + + hunk_sizes = (20, 40, 60, 80) + hunks = [ + "@@ -1 +1 @@\n-old\n+" + ("alpha " * size) + for size in hunk_sizes + ] + files = [ + FilePatchInfo("old\n", "new\n", hunks[index], f"file_{index}.py", edit_type=EDIT_TYPE.MODIFIED) + for index in range(4) + ] + provider = FakeProvider(files) + token_handler = FakeTokenHandler(prompt_tokens=100) + + monkeypatch.setattr(pr_processing, "sort_files_by_main_languages", lambda languages, files: [{"files": files}]) + monkeypatch.setattr(pr_processing, "get_max_tokens", lambda model: 1_700) + + try: + prepared = pr_processing.get_pr_diff( + provider, + token_handler, + "tiny-model", + add_line_numbers_to_hunks=True, + return_remaining_files=True, + return_prepared=True, + ) + + assert isinstance(prepared, pr_processing.PreparedPRDiff) + assert prepared.file_dict + expected_order = [f"file_{index}.py" for index in range(3, -1, -1)] + assert list(prepared.file_dict) == expected_order + calls_after_prepare = token_handler.count_calls + prepared_chunks = pr_processing.get_pr_multi_diffs( + provider, + token_handler, + "tiny-model", + max_calls=3, + add_line_numbers=True, + return_remaining_files=True, + prepared_diff=prepared, + ) + + fresh_provider = FakeProvider([ + FilePatchInfo("old\n", "new\n", hunks[index], f"file_{index}.py", edit_type=EDIT_TYPE.MODIFIED) + for index in range(4) + ]) + fresh_chunks = pr_processing.get_pr_multi_diffs( + fresh_provider, + FakeTokenHandler(prompt_tokens=100), + "tiny-model", + max_calls=3, + add_line_numbers=True, + return_remaining_files=True, + ) + + assert prepared_chunks == fresh_chunks + prepared_diff_list, _ = prepared_chunks + combined_chunks = "\n".join(prepared_diff_list) + assert [combined_chunks.index(filename) for filename in expected_order] == sorted( + combined_chunks.index(filename) for filename in expected_order + ) + assert token_handler.count_calls == calls_after_prepare + assert (provider.diff_calls, provider.language_calls) == (1, 1) + finally: + for key, value in original.items(): + setattr(settings.config, key, value) + + +@pytest.mark.parametrize( + ("prepared_model", "requested_model", "prepared_line_numbers", "requested_line_numbers"), + [ + ("tiny-model", "other-model", True, True), + ("tiny-model", "tiny-model", False, False), + ], +) +def test_prepared_pr_diff_is_not_reused_across_model_or_patch_format( + monkeypatch, prepared_model, requested_model, prepared_line_numbers, requested_line_numbers +): + settings = get_settings() + original_verbosity_level = settings.config.verbosity_level + settings.config.verbosity_level = 0 + hunk = "@@ -1 +1 @@\n-old\n+" + ("alpha " * 60) + files = [ + FilePatchInfo("old\n", "new\n", hunk, f"file_{index}.py", edit_type=EDIT_TYPE.MODIFIED) + for index in range(4) + ] + provider = FakeProvider(files) + token_handler = FakeTokenHandler(prompt_tokens=100) + monkeypatch.setattr(pr_processing, "sort_files_by_main_languages", lambda languages, files: [{"files": files}]) + monkeypatch.setattr(pr_processing, "get_max_tokens", lambda model: 1_700) + + try: + prepared = pr_processing.get_pr_diff( + provider, + token_handler, + prepared_model, + add_line_numbers_to_hunks=prepared_line_numbers, + return_remaining_files=True, + return_prepared=True, + ) + assert isinstance(prepared, pr_processing.PreparedPRDiff) + + pr_processing.get_pr_multi_diffs( + provider, + token_handler, + requested_model, + max_calls=3, + add_line_numbers=requested_line_numbers, + return_remaining_files=True, + prepared_diff=prepared, + ) + + assert (provider.diff_calls, provider.language_calls) == (2, 2) + finally: + settings.config.verbosity_level = original_verbosity_level + + @pytest.mark.parametrize( "call_diff", [ diff --git a/tests/unittest/test_pr_reviewer_core.py b/tests/unittest/test_pr_reviewer_core.py index 985b0b64c7..9068d80e75 100644 --- a/tests/unittest/test_pr_reviewer_core.py +++ b/tests/unittest/test_pr_reviewer_core.py @@ -14,6 +14,8 @@ from pr_agent.git_providers.azuredevops_provider import AzureDevopsProvider from pr_agent.tools.pr_reviewer import PRReviewer, _review_failure_comment +_VALID_PREDICTION = "review:\n summary: prediction" + def _make_reviewer(git_provider=None): reviewer = PRReviewer.__new__(PRReviewer) @@ -87,7 +89,7 @@ def test_review_failure_comment_treats_quoted_false_as_disabled(): @pytest.mark.asyncio async def test_prepare_prediction_requests_remaining_files_and_preserves_tuple_result(): reviewer = _make_prediction_reviewer() - reviewer._get_prediction = AsyncMock(return_value="prediction") + reviewer._get_prediction = AsyncMock(return_value=_VALID_PREDICTION) with patch( "pr_agent.tools.pr_reviewer.get_pr_diff", @@ -105,34 +107,34 @@ async def test_prepare_prediction_requests_remaining_files_and_preserves_tuple_r ) assert reviewer.patches_diff == "diff" assert reviewer.remaining_files_list == ["src/one.py", "docs/two.md"] - assert reviewer.prediction == "prediction" + assert reviewer.prediction == _VALID_PREDICTION @pytest.mark.asyncio async def test_prepare_prediction_accepts_full_diff_string_when_token_budget_is_sufficient(): reviewer = _make_prediction_reviewer() - reviewer._get_prediction = AsyncMock(return_value="prediction") + reviewer._get_prediction = AsyncMock(return_value=_VALID_PREDICTION) with patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value="diff"): await reviewer._prepare_prediction("model") assert reviewer.patches_diff == "diff" assert reviewer.remaining_files_list == [] - assert reviewer.prediction == "prediction" + assert reviewer.prediction == _VALID_PREDICTION @pytest.mark.asyncio async def test_prepare_prediction_keeps_incremental_review_compatible_with_tuple_result(): reviewer = _make_prediction_reviewer() reviewer.incremental = SimpleNamespace(is_incremental=True) - reviewer._get_prediction = AsyncMock(return_value="prediction") + reviewer._get_prediction = AsyncMock(return_value=_VALID_PREDICTION) with patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value=("diff", ["skipped.py"])): await reviewer._prepare_prediction("model") assert reviewer.patches_diff == "diff" assert reviewer.remaining_files_list == ["skipped.py"] - assert reviewer.prediction == "prediction" + assert reviewer.prediction == _VALID_PREDICTION def _render_review(reviewer, remaining_files, supports_gfm_markdown=False): diff --git a/tests/unittest/test_pr_reviewer_fallback_validation.py b/tests/unittest/test_pr_reviewer_fallback_validation.py new file mode 100644 index 0000000000..b4b4ccd877 --- /dev/null +++ b/tests/unittest/test_pr_reviewer_fallback_validation.py @@ -0,0 +1,129 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, call, patch + +import pytest + +from pr_agent.algo.pr_processing import retry_with_fallback_models +from pr_agent.config_loader import get_settings +from pr_agent.tools.pr_reviewer import PRReviewer +from tests.unittest._settings_helpers import restore_settings, snapshot_settings + +_TRACKED_KEYS = ( + "config.model", + "config.fallback_models", + "config.publish_output", + "config.is_auto_command", + "config.propagate_tool_errors", + "openai.deployment_id", + "openai.fallback_deployments", + "pr_reviewer.persistent_comment", +) + +_VALID_REVIEW = """review: + summary: fallback review +""" + + +@pytest.fixture +def fallback_models(): + snapshot = snapshot_settings(_TRACKED_KEYS) + settings = get_settings() + settings.set("config.model", "primary-model") + settings.set("config.fallback_models", ["fallback-model"]) + settings.set("openai.deployment_id", None) + settings.set("openai.fallback_deployments", []) + yield + restore_settings(snapshot) + + +def _make_reviewer(git_provider=None): + reviewer = PRReviewer.__new__(PRReviewer) + reviewer.git_provider = git_provider or MagicMock() + reviewer.token_handler = MagicMock() + reviewer.pr_url = "https://example/pr/1" + reviewer.incremental = SimpleNamespace(is_incremental=False) + reviewer.remaining_files_list = [] + reviewer.prediction = None + reviewer.vars = {} + return reviewer + + +@pytest.mark.asyncio +async def test_malformed_primary_review_uses_fallback_model(fallback_models): + reviewer = _make_reviewer() + reviewer._get_prediction = AsyncMock(side_effect=["not valid review yaml", _VALID_REVIEW]) + + with patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value="diff"): + await retry_with_fallback_models(reviewer._prepare_prediction, git_provider=reviewer.git_provider) + + assert reviewer._get_prediction.await_args_list == [call("primary-model"), call("fallback-model")] + assert reviewer.prediction == _VALID_REVIEW + + +@pytest.mark.asyncio +async def test_malformed_primary_then_valid_fallback_publishes_review_without_failure( + monkeypatch, fallback_models +): + progress_comment = MagicMock() + git_provider = MagicMock() + git_provider.get_files.return_value = ["app.py"] + git_provider.publish_comment.return_value = progress_comment + git_provider.should_publish_review_as_thread.return_value = False + git_provider.supports_review_comment_identity.return_value = False + reviewer = _make_reviewer(git_provider) + reviewer._get_prediction = AsyncMock(side_effect=["not yaml", _VALID_REVIEW]) + reviewer._prepare_pr_review = MagicMock(return_value="rendered fallback review") + + settings = get_settings() + settings.set("config.publish_output", True) + settings.set("config.is_auto_command", False) + settings.set("config.propagate_tool_errors", False) + settings.set("pr_reviewer.persistent_comment", False) + + monkeypatch.setattr("pr_agent.tools.pr_reviewer.extract_and_cache_pr_tickets", AsyncMock()) + monkeypatch.setattr( + "pr_agent.tools.pr_reviewer.fit_related_tickets_to_prompt_budget", + lambda _pr, raw_vars, _system, _user, _model: (raw_vars, reviewer.token_handler), + ) + + with patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value="diff"): + await reviewer.run() + + assert reviewer._get_prediction.await_args_list == [call("primary-model"), call("fallback-model")] + assert git_provider.publish_comment.call_args_list == [ + call("Preparing review...", is_temporary=True), + call("rendered fallback review"), + ] + assert all("Failed to review PR" not in str(args) for args in git_provider.publish_comment.call_args_list) + git_provider.remove_comment.assert_called_once_with(progress_comment) + + +@pytest.mark.asyncio +async def test_all_malformed_models_publish_one_failure_result(monkeypatch, fallback_models): + progress_comment = MagicMock() + git_provider = MagicMock() + git_provider.get_files.return_value = ["app.py"] + git_provider.publish_comment.return_value = progress_comment + reviewer = _make_reviewer(git_provider) + reviewer._get_prediction = AsyncMock(side_effect=["not yaml", "review: {}"]) + + settings = get_settings() + settings.set("config.publish_output", True) + settings.set("config.is_auto_command", False) + settings.set("config.propagate_tool_errors", False) + + monkeypatch.setattr("pr_agent.tools.pr_reviewer.extract_and_cache_pr_tickets", AsyncMock()) + monkeypatch.setattr( + "pr_agent.tools.pr_reviewer.fit_related_tickets_to_prompt_budget", + lambda _pr, raw_vars, _system, _user, _model: (raw_vars, reviewer.token_handler), + ) + + with patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value="diff"): + await reviewer.run() + + assert reviewer._get_prediction.await_args_list == [call("primary-model"), call("fallback-model")] + assert git_provider.publish_comment.call_args_list == [ + call("Preparing review...", is_temporary=True), + call("Failed to review PR"), + ] + git_provider.remove_comment.assert_called_once_with(progress_comment) diff --git a/tests/unittest/test_pr_reviewer_finding_state.py b/tests/unittest/test_pr_reviewer_finding_state.py index b9468c2b9d..34a211bba7 100644 --- a/tests/unittest/test_pr_reviewer_finding_state.py +++ b/tests/unittest/test_pr_reviewer_finding_state.py @@ -5,6 +5,8 @@ import pytest from pr_agent.algo.review_finding_state import ( + _render_resolved_section, + normalize_finding, parse_review_state, reconcile_review_findings, serialize_review_state, @@ -20,10 +22,12 @@ from pr_agent.git_providers.azuredevops_provider import AzureDevopsProvider from pr_agent.git_providers.bitbucket_provider import BitbucketProvider from pr_agent.git_providers.bitbucket_server_provider import BitbucketServerProvider +from pr_agent.git_providers.codecommit_provider import CodeCommitProvider from pr_agent.git_providers.git_provider import GitProvider from pr_agent.git_providers.gitea_provider import GiteaProvider from pr_agent.git_providers.github_provider import GithubProvider from pr_agent.git_providers.gitlab_provider import GitLabProvider +from pr_agent.git_providers.local_git_provider import LocalGitProvider from pr_agent.tools.pr_reviewer import PRReviewer @@ -59,6 +63,54 @@ def _settings(monkeypatch): return settings +def test_normalize_finding_preserves_markdown_while_fingerprint_ignores_rewrapping(): + body = "The loop never ends:\n```text\nwhile attempts < 3:\n attempts += 1\n```" + rewrapped = "The loop never ends: ```text while attempts < 3: attempts += 1 ```" + + normalized = normalize_finding(_finding(body)) + rewrapped_normalized = normalize_finding(_finding(rewrapped)) + + assert normalized["body"] == body + assert normalized["finding_id"] == rewrapped_normalized["finding_id"] + + +def test_resolved_section_preserves_finding_markdown_structure(): + issue = { + "relevant_file": "src/app.py", + "issue_header": "Possible Bug", + "issue_content": "The loop never ends:\n```suggestion\nwhile attempts < 3:\n attempts += 1\n```", + "start_line": 3, + "end_line": 4, + } + finding = PRReviewer._review_finding_from_issue(issue) + active = reconcile_review_findings( + None, + [finding], + allow_resolution=False, + head_sha="head-1", + timestamp="2026-01-01T00:00:00Z", + ).state + resolved = reconcile_review_findings( + active, + [], + allow_resolution=True, + head_sha="head-2", + timestamp="2026-01-01T00:01:00Z", + ).state + + rendered = _render_resolved_section(resolved) + + assert ( + "### src/app.py:3-4\n\n" + "**Possible Issue**\n\n" + "The loop never ends:\n" + "```text\n" + "while attempts < 3:\n" + " attempts += 1\n" + "```" + ) in rendered + + def test_review_finding_state_is_disabled_without_a_provider(monkeypatch): _settings(monkeypatch) @@ -1049,20 +1101,26 @@ async def fake_retry(prepare_fn, model_type=None, git_provider=None): GiteaProvider, BitbucketProvider, BitbucketServerProvider, + CodeCommitProvider, ], - ids=["github", "gitlab", "azure", "gitea", "bitbucket", "bitbucket-server"], + ids=["github", "gitlab", "azure", "gitea", "bitbucket", "bitbucket-server", "codecommit"], ) -def test_legacy_persistent_publish_overrides_accept_shared_arguments(provider_class): +def test_persistent_publish_signatures_accept_shared_arguments(provider_class): parameters = inspect.signature( provider_class.publish_persistent_comment ).parameters + assert "as_thread" in parameters assert "identity_marker" in parameters assert "legacy_initial_header" in parameters assert "require_agent_authorship" not in parameters assert "fallback_on_error" not in parameters +def test_providers_without_override_inherit_persistent_comment_implementation(): + assert LocalGitProvider.publish_persistent_comment is GitProvider.publish_persistent_comment + + def test_oversized_state_degradation_is_safe_on_the_next_run(monkeypatch): settings = _settings(monkeypatch) monkeypatch.setattr(settings.pr_reviewer, "num_max_findings", 100) diff --git a/tests/unittest/test_pr_reviewer_prompt_contract.py b/tests/unittest/test_pr_reviewer_prompt_contract.py index b4d2fdc824..fba8d4ebeb 100644 --- a/tests/unittest/test_pr_reviewer_prompt_contract.py +++ b/tests/unittest/test_pr_reviewer_prompt_contract.py @@ -55,6 +55,42 @@ def _referenced_variables(half): return meta.find_undeclared_variables(environment.parse(template)) +@pytest.mark.parametrize( + ("duplicate_prompt_examples", "expected_user_examples"), + [(False, 0), (True, 1)], +) +def test_ticket_compliance_examples_use_consumer_field( + monkeypatch, + duplicate_prompt_examples, + expected_user_examples, +): + reviewer = _build_reviewer(monkeypatch) + reviewer.vars["related_tickets"] = [ + SimpleNamespace( + ticket_url="https://tracker.example/tickets/1", + title="Ticket", + labels="", + body="", + requirements="", + ) + ] + reviewer.vars["duplicate_prompt_examples"] = duplicate_prompt_examples + + environment = Environment( + autoescape=select_autoescape(default_for_string=False), + undefined=StrictUndefined, + ) + system_prompt = environment.from_string(get_settings().pr_review_prompt.system).render(reviewer.vars) + user_prompt = environment.from_string(get_settings().pr_review_prompt.user).render(reviewer.vars) + + canonical_example = "\n requires_further_human_verification: |" + legacy_example = "\n overall_compliance_level: |" + assert system_prompt.count(canonical_example) == 1 + assert user_prompt.count(canonical_example) == expected_user_examples + assert legacy_example not in system_prompt + assert legacy_example not in user_prompt + + @pytest.mark.parametrize("half", ["system", "user"]) def test_pr_review_prompt_variables_are_all_supplied(monkeypatch, half): reviewer = _build_reviewer(monkeypatch) diff --git a/tests/unittest/test_pr_reviewer_prompt_example_types.py b/tests/unittest/test_pr_reviewer_prompt_example_types.py new file mode 100644 index 0000000000..b22967075b --- /dev/null +++ b/tests/unittest/test_pr_reviewer_prompt_example_types.py @@ -0,0 +1,119 @@ +"""The reviewer prompt's examples must agree with the types its schema declares. + +Named for the example TYPES, not `prompt_contract`: that name is taken by the +Jinja-variable contract tests beside this file, which answer a different +question about the same prompt. + +Issue #3354. The prompts describe the model's output as `class X(BaseModel)` +blocks that exist nowhere as Python, so nothing checks that the schema, the +example and the parser agree — and two fields had drifted in opposite +directions: + + estimated_effort_to_review_[1-5]: int schema + estimated_effort_to_review_[1-5]: | example — a block scalar, i.e. a str + 3 + score: str schema + score: 89 example — a bare int + +Same shape as #3319/#3325 in `pr_add_docs.toml`, and the same test shape: +parse the example the model is shown and hold it to the declared type. This +covers ONE file; #3354's full ask is a real pydantic model per top-level class +with a fixture per tool, which is a larger piece of work. + +Both example blocks are checked, not just the first: they drifted identically, +so a test that read one would go green while the other still taught the wrong +type. +""" + +import re + +import yaml + +from pr_agent.config_loader import get_settings + +# The declared type can be `int`, `List[Foo]` or `Union[List[Foo], str]`, and +# `Field(` may be called positionally, so match up to the `= Field(`. +_DECLARED = re.compile( + r"^\s*(?P[\w\[\]\-]+):\s*(?P[^=]+?)\s*=\s*Field\(", re.M +) + + +def _prompt() -> str: + """Both halves. The schema and the first example live in `system`; the + SECOND example lives in `user`, and they drifted identically -- a test that + read only one would go green while the other still taught the wrong type.""" + prompt = get_settings().pr_review_prompt + return f"{prompt.system}\n{prompt.user}" + + +def _declared_types() -> dict[str, str]: + """field name -> the type its `Field(...)` line declares.""" + return {m.group("field"): m.group("type") for m in _DECLARED.finditer(_prompt())} + + +def _example_blocks() -> list[str]: + """Every ```yaml example in the prompt, Jinja lines stripped. + + The blocks are Jinja-conditional, so the raw text is not valid YAML; the + `{%- ... %}` lines are control flow around the fields, never part of one. + """ + blocks = re.findall(r"```yaml\n(.*?)```", _prompt(), re.S) + return [ + "\n".join(line for line in b.splitlines() if not line.strip().startswith("{%")) + for b in blocks + ] + + +def _review_examples() -> list[dict]: + out = [] + for block in _example_blocks(): + parsed = yaml.safe_load(block) + if isinstance(parsed, dict) and "review" in parsed: + out.append(parsed["review"]) + return out + + +def test_the_prompt_still_carries_the_example_blocks_this_pins(): + """The control: if the prompt stops shipping examples, every assertion + below would pass by having nothing to check.""" + examples = _review_examples() + assert len(examples) == 2, f"expected both review examples, got {len(examples)}" + + +def test_the_effort_example_is_the_int_its_schema_declares(): + assert _declared_types()["estimated_effort_to_review_[1-5]"] == "int" + for example in _review_examples(): + value = example["estimated_effort_to_review_[1-5]"] + assert isinstance(value, int), ( + f"the schema declares int; the example teaches {type(value).__name__} " + f"({value!r}) — a block scalar is a str" + ) + + +def test_the_score_example_is_the_int_its_schema_declares(): + assert _declared_types()["score"] == "int" + for example in _review_examples(): + value = example["score"] + assert isinstance(value, int), ( + f"the schema declares int; the example teaches {type(value).__name__} " + f"({value!r})" + ) + + +def _taught_fields(value) -> set[str]: + """Every key an example teaches, at any depth.""" + if isinstance(value, dict): + return set(value) | {f for v in value.values() for f in _taught_fields(v)} + if isinstance(value, list): + return {f for v in value for f in _taught_fields(v)} + return set() + + +def test_every_field_the_examples_teach_is_one_the_schema_declares(): + """The drift that started #3354 was a field the examples kept after the + schema dropped it (`overall_compliance_level`, #3318). It sat under + `ticket_compliance_check`, so walk nested mappings and lists too.""" + declared = set(_declared_types()) + for example in _review_examples(): + unknown = sorted(_taught_fields(example) - declared) + assert not unknown, f"examples teach fields the schema does not declare: {unknown}" diff --git a/tests/unittest/test_pr_update_changelog.py b/tests/unittest/test_pr_update_changelog.py index b80f15a07f..d7cf0ca878 100644 --- a/tests/unittest/test_pr_update_changelog.py +++ b/tests/unittest/test_pr_update_changelog.py @@ -2,9 +2,17 @@ import pytest +from pr_agent.git_providers.github_provider import GithubProvider from pr_agent.tools.pr_update_changelog import PRUpdateChangelog +class _CustomProvider: + """Provider that opts in without inheriting from GithubProvider.""" + + def supports_changelog_update_review(self) -> bool: + return True + + class TestPRUpdateChangelog: """Test suite for the PR Update Changelog functionality.""" @@ -141,7 +149,7 @@ def test_prepare_changelog_update_no_commit(self, changelog_tool): def _make_no_push_provider(self, extra_spec=None): spec = ["publish_comment", "remove_initial_comment", "get_pr_branch", "get_pr_description", "get_commit_messages", "get_languages", "get_files", "get_pr_file_content", - "is_supported", "pr"] + "is_supported", "supports_changelog_update_review", "pr"] if extra_spec: spec += extra_spec provider = MagicMock(spec=spec) @@ -153,6 +161,7 @@ def _make_no_push_provider(self, extra_spec=None): provider.get_languages.return_value = {"Python": 80, "JavaScript": 20} provider.get_files.return_value = ["test.py", "test.js"] provider.get_pr_file_content.return_value = "" + provider.supports_changelog_update_review.return_value = False return provider @pytest.mark.asyncio @@ -243,6 +252,84 @@ async def test_run_with_push_support(self, changelog_tool, mock_git_provider): assert call_args[1]['file_path'] == 'CHANGELOG.md' assert call_args[1]['branch'] == 'feature-branch' + def test_push_changelog_update_creates_review_when_supported(self, changelog_tool, mock_git_provider): + """When supported, pushing the changelog creates a PR review on the committed changes.""" + mock_git_provider.create_or_update_pr_file = MagicMock() + mock_git_provider.get_pr_branch.return_value = "feature-branch" + mock_git_provider.supports_changelog_update_review.return_value = True + mock_git_provider.pr.get_commits.return_value = ["commit-123"] + mock_git_provider.pr.create_review = MagicMock() + new_content = "# Updated changelog content" + answer = "Line 1\nLine 2" + + with patch("pr_agent.tools.pr_update_changelog.get_settings") as mock_settings, patch( + "pr_agent.tools.pr_update_changelog.sleep" + ): + mock_settings.return_value.pr_update_changelog.get.return_value = True + + changelog_tool._push_changelog_update(new_content, answer) + + mock_git_provider.create_or_update_pr_file.assert_called_once_with( + file_path="CHANGELOG.md", + branch="feature-branch", + contents=new_content, + message="[skip ci] Update CHANGELOG.md", + ) + mock_git_provider.pr.create_review.assert_called_once_with( + commit="commit-123", + comments=[ + dict( + body="CHANGELOG.md update", + path="CHANGELOG.md", + line=2, + start_line=1, + ) + ], + ) + mock_git_provider.publish_comment.assert_not_called() + + def test_push_changelog_update_skips_review_when_not_supported(self, changelog_tool, mock_git_provider): + """A provider without the capability is never asked for a commit-scoped review.""" + mock_git_provider.create_or_update_pr_file = MagicMock() + mock_git_provider.get_pr_branch.return_value = "feature-branch" + mock_git_provider.supports_changelog_update_review.return_value = False + new_content = "# Updated changelog content" + answer = "Changes made" + + with patch("pr_agent.tools.pr_update_changelog.get_settings") as mock_settings, patch( + "pr_agent.tools.pr_update_changelog.sleep" + ): + mock_settings.return_value.pr_update_changelog.get.return_value = True + + changelog_tool._push_changelog_update(new_content, answer) + + mock_git_provider.create_or_update_pr_file.assert_called_once_with( + file_path="CHANGELOG.md", + branch="feature-branch", + contents=new_content, + message="[skip ci] Update CHANGELOG.md", + ) + mock_git_provider.pr.get_commits.assert_not_called() + mock_git_provider.publish_comment.assert_not_called() + + def test_push_changelog_update_falls_back_to_comment_on_review_exception(self, changelog_tool, mock_git_provider): + """When creating a review raises an exception, it falls back to publishing a comment.""" + mock_git_provider.create_or_update_pr_file = MagicMock() + mock_git_provider.get_pr_branch.return_value = "feature-branch" + mock_git_provider.supports_changelog_update_review.return_value = True + mock_git_provider.pr.get_commits.side_effect = Exception("API error") + new_content = "# Updated changelog content" + answer = "Changes made" + + with patch("pr_agent.tools.pr_update_changelog.get_settings") as mock_settings, patch( + "pr_agent.tools.pr_update_changelog.sleep" + ): + mock_settings.return_value.pr_update_changelog.get.return_value = True + + changelog_tool._push_changelog_update(new_content, answer) + + mock_git_provider.publish_comment.assert_called_once_with(f"**Changelog updates: 🔄**\n\n{answer}") + def test_push_changelog_update(self, changelog_tool, mock_git_provider): """Test the push changelog update functionality.""" # Arrange @@ -267,6 +354,20 @@ def test_push_changelog_update(self, changelog_tool, mock_git_provider): message="[skip ci] Update CHANGELOG.md" ) + def test_push_changelog_update_never_calls_create_or_update_pr_file_when_push_code_is_unsupported( + self, changelog_tool, mock_git_provider): + """A provider that declines `push_code` (e.g. restricted_mode) must never reach + `create_or_update_pr_file`, even if a future caller invokes this method directly + without going through `run()`'s own `self.commit_changelog` gate.""" + mock_git_provider.create_or_update_pr_file = MagicMock() + mock_git_provider.is_supported.return_value = False + new_content = "# Updated changelog content" + answer = "Changes made" + + changelog_tool._push_changelog_update(new_content, answer) + + mock_git_provider.create_or_update_pr_file.assert_not_called() + def test_gitlab_provider_method_detection(self, changelog_tool, mock_git_provider): """Test that the tool correctly detects GitLab provider method availability.""" # Arrange @@ -275,6 +376,18 @@ def test_gitlab_provider_method_detection(self, changelog_tool, mock_git_provide # Act & Assert assert hasattr(mock_git_provider, "create_or_update_pr_file") + @pytest.mark.parametrize( + "provider_class, expected", + [(GithubProvider, True), (_CustomProvider, True), (MagicMock, False)], + ) + def test_supports_changelog_update_review_follows_provider_capability(self, provider_class, expected): + if provider_class is MagicMock: + provider = MagicMock() + provider.supports_changelog_update_review.return_value = False + else: + provider = provider_class.__new__(provider_class) + assert provider.supports_changelog_update_review() is expected + @pytest.mark.parametrize("existing_content,new_entry,expected_order", [ ( "# Changelog\n\n## v1.0.0\n- Old feature", diff --git a/tests/unittest/test_publish_workflow_security.py b/tests/unittest/test_publish_workflow_security.py index 5cdd696f0b..b97fb46e9f 100644 --- a/tests/unittest/test_publish_workflow_security.py +++ b/tests/unittest/test_publish_workflow_security.py @@ -185,3 +185,19 @@ def test_release_version_cannot_inject_workflow_outputs(tmp_path: Path, line_bre assert result.returncode != 0 assert output.read_text() == "existing-output\n" + + +def test_release_finalize_keeps_uv_lock_in_sync() -> None: + workflow = yaml.safe_load(PUBLISH_WORKFLOW.read_text()) + finalize_steps = workflow["jobs"]["finalize"]["steps"] + bump_step = next( + step for step in finalize_steps if step.get("name") == "Bump pyproject.toml and sync uv.lock" + ) + commit_step = next( + step for step in finalize_steps if step.get("name") == "Commit & push bump (fast-forward only)" + ) + + run = bump_step["run"] + assert run.index("scripts/set_pyproject_version.py") < run.index("uv lock") + assert "git diff --quiet pyproject.toml uv.lock" in commit_step["run"] + assert "git add pyproject.toml uv.lock" in commit_step["run"] diff --git a/tests/unittest/test_push_outputs.py b/tests/unittest/test_push_outputs.py index dac6d1103e..bd9d8d9450 100644 --- a/tests/unittest/test_push_outputs.py +++ b/tests/unittest/test_push_outputs.py @@ -1,4 +1,5 @@ import json +from types import SimpleNamespace import pytest @@ -14,6 +15,7 @@ def _reset_push_outputs(): s = get_settings() s.set('PUSH_OUTPUTS.ENABLE', False) s.set('PUSH_OUTPUTS.CHANNELS', []) + s.set('PUSH_OUTPUTS.FILE_PATH', 'pr-agent-outputs/reviews.jsonl') s.set('PUSH_OUTPUTS.WEBHOOK_URL', '') s.set('PUSH_OUTPUTS.SLACK_WEBHOOK_URL', '') @@ -65,6 +67,7 @@ def test_slack_channel_posts_text(self, monkeypatch): def fake_post(url, json=None, timeout=None, **kwargs): captured['url'] = url captured['json'] = json + return SimpleNamespace(status_code=200) monkeypatch.setattr(utils.requests, 'post', fake_post) @@ -103,7 +106,7 @@ def get_repo_settings(self): def test_errors_are_non_fatal(self, monkeypatch): get_settings().set('PUSH_OUTPUTS.ENABLE', True) get_settings().set('PUSH_OUTPUTS.CHANNELS', ['webhook']) - get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', 'http://example.invalid/hook') + get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', 'https://example.invalid/hook') def boom(*args, **kwargs): raise ConnectionError("no network") @@ -128,7 +131,7 @@ def test_non_https_sink_urls_are_ignored(self, monkeypatch, bad_url): posts = [] monkeypatch.setattr(utils.requests, 'post', - lambda url, **kwargs: posts.append(url)) + lambda url, **kwargs: posts.append(url) or SimpleNamespace(status_code=200)) push_outputs("review", payload={"a": 1}, markdown="hi") @@ -141,8 +144,210 @@ def test_https_sink_url_is_accepted(self, monkeypatch): posts = [] monkeypatch.setattr(utils.requests, 'post', - lambda url, **kwargs: posts.append(url)) + lambda url, **kwargs: posts.append(url) or SimpleNamespace(status_code=200)) push_outputs("review", payload={"a": 1}, markdown="hi") assert posts == ['https://example.test/hook'] + + def test_setup_errors_remain_non_fatal_and_secret_safe(self, monkeypatch): + warnings = [] + + def fail_settings(): + raise RuntimeError("secret setup marker") + + monkeypatch.setattr(utils, 'get_settings', fail_settings) + monkeypatch.setattr(utils, 'get_logger', + lambda: SimpleNamespace(warning=warnings.append)) + + push_outputs("review", payload={"payload-secret": 1}, markdown="markdown-secret") + + assert warnings == ["push_outputs failed: RuntimeError"] + assert "secret" not in warnings[0] + + def test_webhook_exception_does_not_skip_slack(self, monkeypatch): + get_settings().set('PUSH_OUTPUTS.ENABLE', True) + get_settings().set('PUSH_OUTPUTS.CHANNELS', ['webhook', 'slack']) + webhook_url = 'https://example.test/webhook-secret' + slack_url = 'https://example.test/slack-secret' + get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', webhook_url) + get_settings().set('PUSH_OUTPUTS.SLACK_WEBHOOK_URL', slack_url) + posts = [] + warnings = [] + + def fake_post(url, **kwargs): + posts.append((url, kwargs)) + if url == webhook_url: + raise ConnectionError("transport-secret") + return SimpleNamespace(status_code=200) + + monkeypatch.setattr(utils.requests, 'post', fake_post) + monkeypatch.setattr(utils, 'get_logger', + lambda: SimpleNamespace(warning=warnings.append)) + + push_outputs("review", payload={"payload-secret": 1}, markdown="markdown-secret") + + assert [url for url, _ in posts] == [webhook_url, slack_url] + assert posts[0][1]['timeout'] == 5 + assert posts[0][1]['allow_redirects'] is False + assert posts[1][1]['timeout'] == 5 + assert posts[1][1]['allow_redirects'] is False + assert warnings == ["push_outputs: webhook failed: ConnectionError"] + assert not any(secret in warnings[0] for secret in + (webhook_url, slack_url, "transport-secret", "payload-secret", "markdown-secret")) + + @pytest.mark.parametrize("status_code", [302, 500]) + def test_webhook_non_2xx_warns_and_does_not_skip_slack(self, monkeypatch, status_code): + get_settings().set('PUSH_OUTPUTS.ENABLE', True) + get_settings().set('PUSH_OUTPUTS.CHANNELS', ['webhook', 'slack']) + webhook_url = 'https://example.test/webhook' + slack_url = 'https://example.test/slack' + get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', webhook_url) + get_settings().set('PUSH_OUTPUTS.SLACK_WEBHOOK_URL', slack_url) + posts = [] + warnings = [] + + def fake_post(url, **kwargs): + posts.append(url) + if url == webhook_url: + return SimpleNamespace(status_code=status_code, text="response-secret") + return SimpleNamespace(status_code=204) + + monkeypatch.setattr(utils.requests, 'post', fake_post) + monkeypatch.setattr(utils, 'get_logger', + lambda: SimpleNamespace(warning=warnings.append)) + + push_outputs("review", payload={"a": 1}, markdown="hi") + + assert posts == [webhook_url, slack_url] + assert warnings == [f"push_outputs: webhook failed with status {status_code}"] + assert "response-secret" not in warnings[0] + + @pytest.mark.parametrize("status_code", [200, 204]) + def test_remote_2xx_responses_are_silent(self, monkeypatch, status_code): + get_settings().set('PUSH_OUTPUTS.ENABLE', True) + get_settings().set('PUSH_OUTPUTS.CHANNELS', ['webhook', 'slack']) + get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', 'https://example.test/webhook') + get_settings().set('PUSH_OUTPUTS.SLACK_WEBHOOK_URL', 'https://example.test/slack') + warnings = [] + monkeypatch.setattr(utils.requests, 'post', + lambda *args, **kwargs: SimpleNamespace(status_code=status_code)) + monkeypatch.setattr(utils, 'get_logger', + lambda: SimpleNamespace(warning=warnings.append)) + + push_outputs("review", payload={"a": 1}, markdown="hi") + + assert warnings == [] + + def test_malformed_webhook_url_does_not_skip_slack(self, monkeypatch): + get_settings().set('PUSH_OUTPUTS.ENABLE', True) + get_settings().set('PUSH_OUTPUTS.CHANNELS', ['webhook', 'slack']) + get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', 'https://[') + slack_url = 'https://example.test/slack' + get_settings().set('PUSH_OUTPUTS.SLACK_WEBHOOK_URL', slack_url) + posts = [] + warnings = [] + monkeypatch.setattr(utils.requests, 'post', + lambda url, **kwargs: posts.append(url) or SimpleNamespace(status_code=200)) + monkeypatch.setattr(utils, 'get_logger', + lambda: SimpleNamespace(warning=warnings.append)) + + push_outputs("review", payload={"a": 1}, markdown="hi") + + assert posts == [slack_url] + assert warnings == ["push_outputs: webhook failed: ValueError"] + assert "https://[" not in warnings[0] + + def test_stdout_failure_does_not_skip_later_destinations(self, monkeypatch, tmp_path): + get_settings().set('PUSH_OUTPUTS.ENABLE', True) + get_settings().set('PUSH_OUTPUTS.CHANNELS', ['stdout', 'file', 'webhook', 'slack']) + out = tmp_path / 'out.jsonl' + get_settings().set('PUSH_OUTPUTS.FILE_PATH', str(out)) + get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', 'https://example.test/webhook') + get_settings().set('PUSH_OUTPUTS.SLACK_WEBHOOK_URL', 'https://example.test/slack') + posts = [] + warnings = [] + + def fail_print(*args, **kwargs): + raise OSError("stdout-secret") + + monkeypatch.setattr('builtins.print', fail_print) + monkeypatch.setattr(utils.requests, 'post', + lambda url, **kwargs: posts.append(url) or SimpleNamespace(status_code=200)) + monkeypatch.setattr(utils, 'get_logger', + lambda: SimpleNamespace(warning=warnings.append)) + + push_outputs("review", payload={"a": 1}, markdown="hi") + + assert out.exists() + assert posts == ['https://example.test/webhook', 'https://example.test/slack'] + assert warnings == ["push_outputs: stdout failed: OSError"] + + def test_file_failure_does_not_skip_remote_destinations(self, monkeypatch, tmp_path): + get_settings().set('PUSH_OUTPUTS.ENABLE', True) + get_settings().set('PUSH_OUTPUTS.CHANNELS', ['file', 'webhook', 'slack']) + get_settings().set('PUSH_OUTPUTS.FILE_PATH', str(tmp_path)) + get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', 'https://example.test/webhook') + get_settings().set('PUSH_OUTPUTS.SLACK_WEBHOOK_URL', 'https://example.test/slack') + posts = [] + warnings = [] + monkeypatch.setattr(utils.requests, 'post', + lambda url, **kwargs: posts.append(url) or SimpleNamespace(status_code=200)) + monkeypatch.setattr(utils, 'get_logger', + lambda: SimpleNamespace(warning=warnings.append)) + + push_outputs("review", payload={"a": 1}, markdown="hi") + + assert posts == ['https://example.test/webhook', 'https://example.test/slack'] + assert warnings == ["push_outputs: file failed: IsADirectoryError"] + + def test_slack_failure_is_non_fatal_and_secret_safe(self, monkeypatch): + get_settings().set('PUSH_OUTPUTS.ENABLE', True) + get_settings().set('PUSH_OUTPUTS.CHANNELS', ['webhook', 'slack']) + get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', 'https://example.test/webhook') + slack_url = 'https://example.test/slack-secret' + get_settings().set('PUSH_OUTPUTS.SLACK_WEBHOOK_URL', slack_url) + warnings = [] + + def fake_post(url, **kwargs): + if url == slack_url: + raise TimeoutError("slack-timeout-secret") + return SimpleNamespace(status_code=200) + + monkeypatch.setattr(utils.requests, 'post', fake_post) + monkeypatch.setattr(utils, 'get_logger', + lambda: SimpleNamespace(warning=warnings.append)) + + push_outputs("review", payload={"a": 1}, markdown="hi") + + assert warnings == ["push_outputs: slack failed: TimeoutError"] + assert slack_url not in warnings[0] + assert "slack-timeout-secret" not in warnings[0] + + def test_slack_non_2xx_warns_without_response_content(self, monkeypatch): + get_settings().set('PUSH_OUTPUTS.ENABLE', True) + get_settings().set('PUSH_OUTPUTS.CHANNELS', ['slack']) + get_settings().set('PUSH_OUTPUTS.SLACK_WEBHOOK_URL', 'https://example.test/slack') + warnings = [] + monkeypatch.setattr(utils.requests, 'post', + lambda *args, **kwargs: SimpleNamespace(status_code=429, text="response-secret")) + monkeypatch.setattr(utils, 'get_logger', + lambda: SimpleNamespace(warning=warnings.append)) + + push_outputs("review", payload={"a": 1}, markdown="hi") + + assert warnings == ["push_outputs: slack failed with status 429"] + assert "response-secret" not in warnings[0] + + def test_unknown_and_duplicate_channels_do_not_duplicate_delivery(self, monkeypatch): + get_settings().set('PUSH_OUTPUTS.ENABLE', True) + get_settings().set('PUSH_OUTPUTS.CHANNELS', ['unknown', 'webhook', 'webhook']) + webhook_url = 'https://example.test/webhook' + get_settings().set('PUSH_OUTPUTS.WEBHOOK_URL', webhook_url) + posts = [] + monkeypatch.setattr(utils.requests, 'post', + lambda url, **kwargs: posts.append(url) or SimpleNamespace(status_code=200)) + + push_outputs("review", payload={"a": 1}, markdown="hi") + + assert posts == [webhook_url] diff --git a/tests/unittest/test_repo_context.py b/tests/unittest/test_repo_context.py index 67075577f1..aac196e5d1 100644 --- a/tests/unittest/test_repo_context.py +++ b/tests/unittest/test_repo_context.py @@ -1,7 +1,9 @@ +from types import SimpleNamespace from unittest.mock import Mock, patch import pytest from github import GithubException +from gitlab.exceptions import GitlabGetError from jinja2 import Environment, StrictUndefined, select_autoescape from pr_agent.algo import repo_context @@ -13,6 +15,14 @@ render_instruction_files_with_line_budget, ) from pr_agent.config_loader import get_settings +from pr_agent.git_providers import ( + AzureDevopsProvider, + BitbucketProvider, + BitbucketServerProvider, + GiteaProvider, + GitLabProvider, +) +from pr_agent.git_providers.codecommit_provider import CodeCommitProvider from pr_agent.git_providers.git_provider import GitProvider from pr_agent.git_providers.github_provider import GithubProvider @@ -29,25 +39,48 @@ def get_repo_file_content(self, file_path: str, from_default_branch: bool = Fals self.from_default_branch_calls.append(from_default_branch) return self.files.get(file_path) + def get_repo_context_ref(self, from_default_branch: bool = False): + # Simulate a provider that keys its content on a revision: the default branch uses a + # stable name, and the target/base branch uses a commit-derived value. + return "default" if from_default_branch else "target-sha" + class UnsupportedProvider: get_repo_file_content = GitProvider.get_repo_file_content +class SiblingFakeProvider(FakeProvider): + """FakeProvider that additionally resolves sibling-repository files.""" + + def __init__(self, files, sibling_files=None, pr_url=None): + super().__init__(files, pr_url) + self.sibling_files = sibling_files or {} + self.requested_siblings = [] + + def get_sibling_repo_file_content(self, repo_id: str, file_path: str, from_default_branch: bool = False): + self.requested_siblings.append(f"{repo_id}:{file_path}") + return self.sibling_files.get(f"{repo_id}:{file_path}") + + @pytest.fixture def repo_context_settings(): settings = get_settings() original_files = settings.config.get("repo_context_files", []) original_max_lines = settings.config.get("repo_context_max_lines", 500) original_from_default_branch = settings.config.get("repo_context_from_default_branch", True) + original_sibling_repos = settings.config.get("repo_context_sibling_repos", []) + settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", []) + original_max_sibling_files = settings.config.get("repo_context_max_sibling_files", 5) original_warned_provider_classes = repo_context._unsupported_repo_context_provider_classes.copy() original_process_cache = repo_context._repo_context_process_cache.copy() yield settings + settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", original_sibling_repos) settings.set("CONFIG.REPO_CONTEXT_FILES", original_files) settings.set("CONFIG.REPO_CONTEXT_MAX_LINES", original_max_lines) settings.set("CONFIG.REPO_CONTEXT_FROM_DEFAULT_BRANCH", original_from_default_branch) + settings.set("CONFIG.REPO_CONTEXT_MAX_SIBLING_FILES", original_max_sibling_files) repo_context._unsupported_repo_context_provider_classes = original_warned_provider_classes repo_context._repo_context_process_cache = original_process_cache @@ -287,7 +320,7 @@ def test_load_repo_context_files_normalizes_fetch_results(): provider, ["AGENTS.md", "EMPTY.md", "MISSING.md", " "] ) - assert files == {"AGENTS.md": "Repo purpose"} + assert files == [("AGENTS.md", "Repo purpose")] assert had_fetch_error is False assert provider.requested_paths == ["AGENTS.md", "EMPTY.md", "MISSING.md"] @@ -298,10 +331,937 @@ def test_load_repo_context_files_reports_fetch_errors(): files, had_fetch_error = repo_context._load_repo_context_files(provider, ["AGENTS.md"]) - assert files == {} + assert files == [] + assert had_fetch_error is True + + +@pytest.mark.parametrize( + "entry,expected_repo_id,expected_path", + [ + ("AGENTS.md", None, "AGENTS.md"), + ("docs/guide.md", None, "docs/guide.md"), + (" src/ x.py ", None, "src/ x.py"), + ("docs/guide:part.md", None, "docs/guide:part.md"), + ({"repo_id": "group/sub/lib", "file_path": "src/api.py"}, "group/sub/lib", "src/api.py"), + ({"repo_id": "owner/lib", "file_path": "/deep/file.py"}, "owner/lib", "deep/file.py"), + ({"repo_id": "owner/lib", "file_path": " file.py "}, "owner/lib", "file.py"), + ({"repo_id": " owner/lib ", "file_path": " file.py "}, "owner/lib", "file.py"), + ], +) +def test_parse_repo_context_file_entry(entry, expected_repo_id, expected_path): + assert repo_context._parse_repo_context_file_entry(entry) == (expected_repo_id, expected_path) + + +@pytest.mark.parametrize( + "entry", + [ + {"repo_id": "group/lib", "file_path": ""}, + {"repo_id": "", "file_path": "src/api.py"}, + {"repo_id": "group/lib"}, + {"file_path": "src/api.py"}, + {"repo_id": 123, "file_path": "src/api.py"}, + {"repo_id": "group/lib", "file_path": "/"}, + 123, + ], +) +def test_parse_repo_context_file_entry_rejects_malformed(entry): + # The parser is pure: malformed entries read as invalid, and the loader reports the skip. + assert repo_context._parse_repo_context_file_entry(entry) == (None, "") + + +def test_parse_repo_context_file_entry_treats_blank_as_plain_path(): + assert repo_context._parse_repo_context_file_entry(" ") == (None, "") + assert repo_context._parse_repo_context_file_entry("") == (None, "") + + +def test_load_repo_context_files_fetches_sibling_and_same_repo_files(): + provider = SiblingFakeProvider( + files={"AGENTS.md": "Repo purpose"}, + sibling_files={"group/lib-api:src/interfaces/api.py": "Sibling contract"}, + ) + + files, had_fetch_error = repo_context._load_repo_context_files( + provider, + ["AGENTS.md", {"repo_id": "group/lib-api", "file_path": "src/interfaces/api.py"}], + from_default_branch=True, + ) + + # The sibling file is rendered under its sibling path so the model sees where it came from. + assert files == [ + ("AGENTS.md", "Repo purpose"), + ("group/lib-api/src/interfaces/api.py", "Sibling contract"), + ] + assert had_fetch_error is False + assert provider.requested_paths == ["AGENTS.md"] + assert provider.requested_siblings == ["group/lib-api:src/interfaces/api.py"] + + +def test_load_repo_context_files_treats_colon_string_as_local_path(): + provider = SiblingFakeProvider( + files={"docs/guide:part.md": "Local content"}, + sibling_files={"docs/guide:part.md": "Sibling content"}, + ) + + files, had_fetch_error = repo_context._load_repo_context_files( + provider, ["docs/guide:part.md"], from_default_branch=True + ) + + # A ':' inside a plain string is a local path, never a sibling reference. + assert files == [("docs/guide:part.md", "Local content")] + assert had_fetch_error is False + assert provider.requested_paths == ["docs/guide:part.md"] + assert provider.requested_siblings == [] + + +def test_load_repo_context_files_reports_sibling_fetch_errors(): + provider = SiblingFakeProvider(files={}) + provider.get_sibling_repo_file_content = Mock(side_effect=Exception("temporary outage")) + + files, had_fetch_error = repo_context._load_repo_context_files( + provider, [{"repo_id": "group/lib-api", "file_path": "src/api.py"}], from_default_branch=True + ) + + assert files == [] + assert had_fetch_error is True + + +def test_load_repo_context_files_skips_siblings_for_unsupported_provider(): + provider = FakeProvider({"AGENTS.md": "Repo purpose"}) + provider.get_sibling_repo_file_content = Mock() + + files, had_fetch_error = repo_context._load_repo_context_files( + provider, ["AGENTS.md", {"repo_id": "group/lib-api", "file_path": "src/api.py"}], from_default_branch=True + ) + + assert files == [("AGENTS.md", "Repo purpose")] + assert had_fetch_error is False + assert provider.requested_paths == ["AGENTS.md"] + provider.get_sibling_repo_file_content.assert_not_called() + + +def test_load_repo_context_files_respects_sibling_file_cap(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_SIBLING_FILES", 2) + provider = SiblingFakeProvider( + files={}, + sibling_files={ + "group/g1:README.md": "one", + "group/g2:README.md": "two", + "group/g3:README.md": "three", + }, + ) + + files, had_fetch_error = repo_context._load_repo_context_files( + provider, + [ + {"repo_id": "group/g1", "file_path": "README.md"}, + {"repo_id": "group/g2", "file_path": "README.md"}, + {"repo_id": "group/g3", "file_path": "README.md"}, + {"repo_id": "group/g4", "file_path": "README.md"}, + ], + from_default_branch=True, + ) + + assert files == [ + ("group/g1/README.md", "one"), + ("group/g2/README.md", "two"), + ] + assert had_fetch_error is False + assert provider.requested_siblings == [ + "group/g1:README.md", + "group/g2:README.md", + ] + + +def test_sibling_fetch_cap_counts_unique_pairs_not_duplicate_entries(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_SIBLING_FILES", 2) + provider = SiblingFakeProvider( + files={}, + sibling_files={ + "group/g1:README.md": "one", + "group/g2:README.md": "two", + }, + ) + + files, had_fetch_error = repo_context._load_repo_context_files( + provider, + [ + {"repo_id": "group/g1", "file_path": "README.md"}, + {"repo_id": "group/g1", "file_path": "README.md"}, + {"repo_id": "group/g2", "file_path": "README.md"}, + ], + from_default_branch=True, + ) + + # A duplicate entry must not consume the fetch cap: each unique pair is fetched once. + assert files == [ + ("group/g1/README.md", "one"), + ("group/g2/README.md", "two"), + ] + assert had_fetch_error is False + assert provider.requested_siblings == [ + "group/g1:README.md", + "group/g2:README.md", + ] + + +def test_read_max_sibling_context_files_clamps_oversized_value(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_SIBLING_FILES", 10**6) + assert repo_context._read_max_sibling_context_files() == repo_context._HARD_MAX_SIBLING_CONTEXT_FILES + + +def test_read_max_sibling_context_files_defaults_on_bad_value(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_SIBLING_FILES", "not-a-number") + assert repo_context._read_max_sibling_context_files() == 5 + + +def test_read_max_sibling_context_files_floor_at_zero(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_SIBLING_FILES", -100) + assert repo_context._read_max_sibling_context_files() == 0 + + +def test_load_repo_context_files_clamps_sibling_fetches_to_hard_ceiling(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_SIBLING_FILES", 10**6) + hard_max = repo_context._HARD_MAX_SIBLING_CONTEXT_FILES + context_files = [{"repo_id": f"group/g{i}", "file_path": "README.md"} for i in range(hard_max + 5)] + provider = SiblingFakeProvider( + files={}, + sibling_files={f"group/g{i}:README.md": str(i) for i in range(hard_max + 5)}, + ) + + files, had_fetch_error = repo_context._load_repo_context_files( + provider, context_files, from_default_branch=True + ) + + assert had_fetch_error is False + assert provider.requested_siblings == [f"group/g{i}:README.md" for i in range(hard_max)] + + +def test_sibling_fetch_cap_counts_attempts_not_just_content(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_SIBLING_FILES", 2) + provider = SiblingFakeProvider( + files={}, + sibling_files={ + "group/g1:README.md": "", + "group/g2:README.md": "two", + }, + ) + + files, had_fetch_error = repo_context._load_repo_context_files( + provider, + [ + {"repo_id": "group/g1", "file_path": "README.md"}, + {"repo_id": "group/g2", "file_path": "README.md"}, + {"repo_id": "group/g3", "file_path": "README.md"}, + ], + from_default_branch=True, + ) + + assert files == [("group/g2/README.md", "two")] + assert provider.requested_siblings == [ + "group/g1:README.md", + "group/g2:README.md", + ] + + +def test_sibling_fetch_cap_counts_exceptions(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_SIBLING_FILES", 2) + provider = SiblingFakeProvider(files={}) + call_count = 0 + original = provider.get_sibling_repo_file_content + + def counting_call(repo_id, file_path): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise Exception("transient outage") + return original(repo_id, file_path) + + provider.get_sibling_repo_file_content = Mock(side_effect=counting_call) + provider.requested_siblings = [] + + files, had_fetch_error = repo_context._load_repo_context_files( + provider, + [ + {"repo_id": "group/g1", "file_path": "README.md"}, + {"repo_id": "group/g2", "file_path": "README.md"}, + {"repo_id": "group/g3", "file_path": "README.md"}, + ], + from_default_branch=True, + ) + + assert call_count == 2 assert had_fetch_error is True +def test_build_repo_context_renders_sibling_file_with_budget(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", ["group/lib-api"]) + repo_context_settings.set( + "CONFIG.REPO_CONTEXT_FILES", + [{"repo_id": "group/lib-api", "file_path": "src/interfaces/api.py"}] + ) + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_LINES", 500) + provider = SiblingFakeProvider( + files={}, + sibling_files={"group/lib-api:src/interfaces/api.py": "def call(req): ...\n"}, + ) + + context = build_repo_context(provider) + + assert context == ( + "You are being given instruction files. Follow them as project-specific guidance when reviewing code.\n" + "\n" + '\n' + "`````markdown\n" + "def call(req): ...\n" + "`````\n" + "\n\n" + "" + ) + + +def test_build_repo_context_renders_local_and_sibling_with_same_label(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", ["group/lib"]) + repo_context_settings.set( + "CONFIG.REPO_CONTEXT_FILES", + [ + "group/lib/README.md", + {"repo_id": "group/lib", "file_path": "README.md"}, + ], + ) + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_LINES", 500) + provider = SiblingFakeProvider( + files={"group/lib/README.md": "Local README\n"}, + sibling_files={"group/lib:README.md": "Sibling README\n"}, + ) + + context = build_repo_context(provider) + + # A local path that equals a sibling's rendered label must not overwrite the other: both + # configured instruction files reach the prompt, distinguished by their source content. + assert context == ( + "You are being given instruction files. Follow them as project-specific guidance when reviewing code.\n" + "\n" + '\n' + "`````markdown\n" + "Local README\n" + "`````\n" + "\n\n" + '\n' + "`````markdown\n" + "Sibling README\n" + "`````\n" + "\n\n" + "" + ) + + +@pytest.fixture +def sibling_provider_settings(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", ["myorg/lib", "group/sub/lib", "456"]) + return repo_context_settings + + +def test_github_provider_fetches_sibling_file_in_same_owner(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = False + sibling_repo.get_contents.return_value.decoded_content = b"sibling contract" + provider.github_client.get_repo.return_value = sibling_repo + + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "sibling contract" + provider.github_client.get_repo.assert_called_once_with("myorg/lib") + sibling_repo.get_contents.assert_called_once_with("src/api.py") + + +def test_github_provider_fetches_private_sibling_when_requester_is_collaborator(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.pr = SimpleNamespace(user=SimpleNamespace(login="alice")) + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = True + sibling_repo.has_in_collaborators.return_value = True + sibling_repo.get_contents.return_value.decoded_content = b"sibling contract" + provider.github_client.get_repo.return_value = sibling_repo + + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "sibling contract" + sibling_repo.has_in_collaborators.assert_called_once_with("alice") + + +def test_github_provider_fetches_private_sibling_when_requester_is_owner(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.pr = SimpleNamespace(user=SimpleNamespace(login="alice")) + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = True + sibling_repo.owner = SimpleNamespace(login="alice") + sibling_repo.get_contents.return_value.decoded_content = b"sibling contract" + provider.github_client.get_repo.return_value = sibling_repo + + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "sibling contract" + sibling_repo.has_in_collaborators.assert_not_called() + + +def test_github_provider_rejects_private_sibling_when_requester_is_not_collaborator(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.pr = SimpleNamespace(user=SimpleNamespace(login="alice")) + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = True + sibling_repo.has_in_collaborators.return_value = False + provider.github_client.get_repo.return_value = sibling_repo + + with patch("pr_agent.git_providers.github_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "" + + sibling_repo.get_contents.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo context file the review requester cannot read: myorg/lib" + ) + + +def test_github_provider_rejects_private_sibling_when_requester_is_unknown(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = True + provider.github_client.get_repo.return_value = sibling_repo + + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "" + sibling_repo.get_contents.assert_not_called() + + +def test_github_provider_rejects_private_sibling_when_command_actor_lacks_access(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.pr = SimpleNamespace(user=SimpleNamespace(login="alice")) + provider.set_command_actor("mallory") + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = True + sibling_repo.visibility = "private" + sibling_repo.has_in_collaborators.return_value = False + provider.github_client.get_repo.return_value = sibling_repo + + with patch("pr_agent.git_providers.github_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "" + + # The commenter (command actor), not the PR author, is the identity that must be verified. + sibling_repo.has_in_collaborators.assert_called_once_with("mallory") + sibling_repo.get_contents.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo context file the review requester cannot read: myorg/lib" + ) + + +def test_github_provider_fetches_private_sibling_when_command_actor_is_collaborator(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.pr = SimpleNamespace(user=SimpleNamespace(login="alice")) + provider.set_command_actor("mallory") + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = True + sibling_repo.visibility = "private" + sibling_repo.has_in_collaborators.return_value = True + sibling_repo.get_contents.return_value.decoded_content = b"sibling contract" + provider.github_client.get_repo.return_value = sibling_repo + + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "sibling contract" + sibling_repo.has_in_collaborators.assert_called_once_with("mallory") + + +def test_github_provider_rejects_internal_sibling_when_requester_is_not_collaborator(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.pr = SimpleNamespace(user=SimpleNamespace(login="alice")) + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = False + sibling_repo.visibility = "internal" + sibling_repo.organization = None # not an organization-controlled repository + sibling_repo.has_in_collaborators.return_value = False + provider.github_client.get_repo.return_value = sibling_repo + + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "" + sibling_repo.has_in_collaborators.assert_called_once_with("alice") + sibling_repo.get_contents.assert_not_called() + + +def test_github_provider_fetches_internal_sibling_when_command_actor_is_collaborator(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.set_command_actor("alice") + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = False + sibling_repo.visibility = "internal" + sibling_repo.organization = None + sibling_repo.has_in_collaborators.return_value = True + sibling_repo.get_contents.return_value.decoded_content = b"sibling contract" + provider.github_client.get_repo.return_value = sibling_repo + + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "sibling contract" + sibling_repo.has_in_collaborators.assert_called_once_with("alice") + + +def test_github_provider_fetches_internal_sibling_when_requester_is_org_member(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.set_command_actor("alice") + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = False + sibling_repo.visibility = "internal" + sibling_repo.organization = Mock() + sibling_repo.organization.has_in_members.return_value = True + sibling_repo.get_contents.return_value.decoded_content = b"sibling contract" + provider.github_client.get_repo.return_value = sibling_repo + + # Internal repositories are readable by every member of the owning organization without a + # per-repo collaborator grant, so the membership check is both sufficient and necessary. + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "sibling contract" + provider.github_client.get_user.assert_called_once_with("alice") + sibling_repo.organization.has_in_members.assert_called_once_with(provider.github_client.get_user.return_value) + sibling_repo.has_in_collaborators.assert_not_called() + + +def test_github_provider_rejects_internal_sibling_to_non_member_non_collaborator(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.set_command_actor("alice") + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = False + sibling_repo.visibility = "internal" + sibling_repo.organization = Mock() + sibling_repo.organization.has_in_members.return_value = False + sibling_repo.has_in_collaborators.return_value = False + provider.github_client.get_repo.return_value = sibling_repo + + with patch("pr_agent.git_providers.github_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "" + + # Not an org member; outside collaborators are the remaining way into an internal repo. + provider.github_client.get_user.assert_called_once_with("alice") + sibling_repo.organization.has_in_members.assert_called_once_with(provider.github_client.get_user.return_value) + sibling_repo.has_in_collaborators.assert_called_once_with("alice") + sibling_repo.get_contents.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo context file the review requester cannot read: myorg/lib" + ) + + +def test_github_provider_fetches_sibling_when_owner_case_differs(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "MyOrg/current" + provider.get_owning_namespace = Mock(return_value="MyOrg") + provider.github_client = Mock() + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = False + sibling_repo.get_contents.return_value.decoded_content = b"sibling contract" + provider.github_client.get_repo.return_value = sibling_repo + + # Owner logins are case-insensitive on GitHub, so differently-cased siblings are valid. + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "sibling contract" + provider.github_client.get_repo.assert_called_once_with("myorg/lib") + sibling_repo.get_contents.assert_called_once_with("src/api.py") + + +def test_github_provider_propagates_sibling_collaborator_errors(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.pr = SimpleNamespace(user=SimpleNamespace(login="alice")) + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = True + sibling_repo.has_in_collaborators.side_effect = GithubException(500, {"message": "boom"}, {}) + provider.github_client.get_repo.return_value = sibling_repo + + # A transient provider failure must surface as a fetch error, not read as a denial. + with pytest.raises(GithubException): + provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") + sibling_repo.get_contents.assert_not_called() + + # A plain denial stays returnable: a definitive 404 means the requester is not a collaborator. + sibling_repo.has_in_collaborators.side_effect = GithubException(404, {"message": "Not Found"}, {}) + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "" + + +def test_github_provider_propagates_org_membership_errors(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.set_command_actor("alice") + sibling_repo = Mock() + sibling_repo.full_name = "myorg/lib" + sibling_repo.private = False + sibling_repo.visibility = "internal" + sibling_repo.organization = Mock() + sibling_repo.organization.has_in_members.side_effect = GithubException(500, {"message": "boom"}, {}) + provider.github_client.get_repo.return_value = sibling_repo + + with pytest.raises(GithubException): + provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") + sibling_repo.has_in_collaborators.assert_not_called() + sibling_repo.get_contents.assert_not_called() + + # A 404 (organization cannot be resolved) is not a membership verdict: fall through. + sibling_repo.organization.has_in_members.side_effect = GithubException(404, {"message": "Not Found"}, {}) + sibling_repo.has_in_collaborators.return_value = False + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "" + provider.github_client.get_user.assert_called_with("alice") + sibling_repo.organization.has_in_members.assert_called_with(provider.github_client.get_user.return_value) + sibling_repo.has_in_collaborators.assert_called_once_with("alice") + + +def test_github_provider_rejects_out_of_owner_sibling_without_api_call(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + + with patch("pr_agent.git_providers.github_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("other-org/lib", "src/api.py") == "" + + provider.github_client.get_repo.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo absent from the host allowlist: other-org/lib" + ) + + +def test_github_provider_treats_missing_sibling_file_as_no_context(sibling_provider_settings): + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + provider.get_owning_namespace = Mock(return_value="myorg") + provider.github_client = Mock() + provider.github_client.get_repo.side_effect = GithubException(404, {"message": "Not Found"}, {}) + + assert provider.get_sibling_repo_file_content("myorg/lib", "src/api.py") == "" + + +def test_gitlab_provider_fetches_sibling_file_in_same_namespace(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "public" + sibling_project.files.get.return_value.decode.return_value = b"sibling contract" + provider.gl.projects.get.return_value = sibling_project + + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "sibling contract" + provider.gl.projects.get.assert_called_once_with("group/sub/lib") + sibling_project.files.get.assert_called_once_with(file_path="src/api.py", ref="main") + + +def test_gitlab_provider_rejects_internal_sibling_when_requester_is_not_member(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "internal" + sibling_project.members_all.get.side_effect = GitlabGetError("Not found", response_code=404) + provider.gl.projects.get.return_value = sibling_project + + with patch("pr_agent.git_providers.gitlab_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "" + + # Internal projects forbid external users and low-role/unknown members, so membership + # must be positively established before the sibling is read. + sibling_project.members_all.get.assert_called_once_with(42) + sibling_project.files.get.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo context file the review requester cannot read: group/sub/lib" + ) + + +def test_gitlab_provider_fetches_internal_sibling_when_requester_is_read_member(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "internal" + sibling_project.members_all.get.return_value = SimpleNamespace(access_level=30) + sibling_project.files.get.return_value.decode.return_value = b"sibling contract" + provider.gl.projects.get.return_value = sibling_project + + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "sibling contract" + sibling_project.members_all.get.assert_called_once_with(42) + + +def test_gitlab_provider_fetches_internal_sibling_for_non_external_instance_user(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + # The requester is a signed-in instance user who is not an external user. + provider.gl.users.get.return_value = SimpleNamespace(external=False) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "internal" + sibling_project.files.get.return_value.decode.return_value = b"sibling contract" + provider.gl.projects.get.return_value = sibling_project + + # Internal projects are visible to every non-external instance user without membership. + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "sibling contract" + provider.gl.users.get.assert_called_once_with(42) + sibling_project.members_all.get.assert_not_called() + sibling_project.files.get.assert_called_once_with(file_path="src/api.py", ref="main") + + +def test_gitlab_provider_rejects_internal_sibling_for_external_user_non_member(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + provider.gl.users.get.return_value = SimpleNamespace(external=True) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "internal" + sibling_project.members_all.get.side_effect = GitlabGetError("Not found", response_code=404) + provider.gl.projects.get.return_value = sibling_project + + with patch("pr_agent.git_providers.gitlab_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "" + + # External users cannot be granted internal access as members either, so this fails closed. + provider.gl.users.get.assert_called_once_with(42) + sibling_project.members_all.get.assert_called_once_with(42) + sibling_project.files.get.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo context file the review requester cannot read: group/sub/lib" + ) + + +def test_gitlab_provider_propagates_membership_provider_errors(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "private" + sibling_project.members_all.get.side_effect = GitlabGetError("boom", response_code=500) + provider.gl.projects.get.return_value = sibling_project + + # A transient membership provider failure must surface as a fetch error, not a denial. + with pytest.raises(GitlabGetError): + provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") + sibling_project.files.get.assert_not_called() + + +def test_gitlab_provider_propagates_user_lookup_provider_errors(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + provider.gl.users.get.side_effect = GitlabGetError("boom", response_code=500) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "internal" + provider.gl.projects.get.return_value = sibling_project + + with pytest.raises(GitlabGetError): + provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") + sibling_project.members_all.get.assert_not_called() + sibling_project.files.get.assert_not_called() + + # A definitive 404 (no such user) is a plain denial: the requester cannot be a member. + provider.gl.users.get.side_effect = GitlabGetError("Not found", response_code=404) + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "" + sibling_project.members_all.get.assert_not_called() + + +def test_gitlab_provider_rejects_private_sibling_when_member_lacks_repo_read(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "private" + sibling_project.members_all.get.return_value = SimpleNamespace(access_level=10) + sibling_project.files.get.return_value.decode.return_value = b"sibling contract" + provider.gl.projects.get.return_value = sibling_project + + with patch("pr_agent.git_providers.gitlab_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "" + + # Guests and minimal-access members cannot read private repository source files. + sibling_project.members_all.get.assert_called_once_with(42) + sibling_project.files.get.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo context file the review requester cannot read: group/sub/lib" + ) + + +def test_gitlab_provider_fetches_private_sibling_when_requester_is_member(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "private" + sibling_project.members_all.get.return_value = SimpleNamespace(access_level=30) + sibling_project.files.get.return_value.decode.return_value = b"sibling contract" + provider.gl.projects.get.return_value = sibling_project + + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "sibling contract" + sibling_project.members_all.get.assert_called_once_with(42) + + +def test_gitlab_provider_rejects_private_sibling_when_command_actor_is_not_member(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + provider.set_command_actor(99) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "private" + sibling_project.members_all.get.side_effect = GitlabGetError("Not found", response_code=404) + provider.gl.projects.get.return_value = sibling_project + + with patch("pr_agent.git_providers.gitlab_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "" + + # The commenter (command actor), not the MR author, is the identity that must be verified. + sibling_project.members_all.get.assert_called_once_with(99) + sibling_project.files.get.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo context file the review requester cannot read: group/sub/lib" + ) + + +def test_gitlab_provider_fetches_private_sibling_when_command_actor_is_member(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + provider.set_command_actor(99) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "private" + sibling_project.members_all.get.return_value = SimpleNamespace(access_level=30) + sibling_project.files.get.return_value.decode.return_value = b"sibling contract" + provider.gl.projects.get.return_value = sibling_project + + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "sibling contract" + sibling_project.members_all.get.assert_called_once_with(99) + + +def test_gitlab_provider_rejects_private_sibling_when_requester_is_not_member(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + provider.mr = SimpleNamespace(author={"id": 42}) + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "private" + sibling_project.members_all.get.side_effect = GitlabGetError("Not found", response_code=404) + provider.gl.projects.get.return_value = sibling_project + + with patch("pr_agent.git_providers.gitlab_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "" + + sibling_project.files.get.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo context file the review requester cannot read: group/sub/lib" + ) + + +def test_gitlab_provider_rejects_private_sibling_when_requester_is_unknown(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + sibling_project = Mock() + sibling_project.path_with_namespace = "group/sub/lib" + sibling_project.default_branch = "main" + sibling_project.visibility = "private" + provider.gl.projects.get.return_value = sibling_project + + assert provider.get_sibling_repo_file_content("group/sub/lib", "src/api.py") == "" + sibling_project.files.get.assert_not_called() + sibling_project.members_all.get.assert_not_called() + + +def test_gitlab_provider_rejects_out_of_namespace_sibling_without_api_call(sibling_provider_settings): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/sub/current" + provider.get_owning_namespace = Mock(return_value="group") + provider.gl = Mock() + + with patch("pr_agent.git_providers.gitlab_provider.get_logger") as mock_get_logger: + assert provider.get_sibling_repo_file_content("other/lib", "src/api.py") == "" + + provider.gl.projects.get.assert_not_called() + mock_get_logger.return_value.warning.assert_called_once_with( + "Ignoring sibling repo absent from the host allowlist: other/lib" + ) + + def test_build_repo_context_process_cache_invalidates_when_config_changes(repo_context_settings): repo_context_settings.set("CONFIG.REPO_CONTEXT_FILES", ["AGENTS.md"]) repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_LINES", 500) @@ -335,7 +1295,130 @@ def test_build_repo_context_does_not_cache_empty_context_after_fetch_error(repo_ assert first_context == "" assert "Repo purpose" in second_context - assert provider.get_repo_file_content.call_count == 2 + + +def test_build_repo_context_does_not_cache_sibling_content(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", ["group/lib"]) + repo_context_settings.set( + "CONFIG.REPO_CONTEXT_FILES", + [{"repo_id": "group/lib", "file_path": "api.py"}] + ) + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_LINES", 500) + provider = SiblingFakeProvider( + files={}, + sibling_files={"group/lib:api.py": "def call(req): ...\n"}, + pr_url="https://example.com/org/repo/pull/1", + ) + + first_context = build_repo_context(provider) + assert "' in rendered + + +class RefishProvider(FakeProvider): + """A provider whose repo-context revision can move between calls, like a rebased base branch.""" + + def __init__(self, files, pr_url=None): + super().__init__(files, pr_url) + self.context_ref = "sha-1" + + def get_repo_context_ref(self, from_default_branch: bool = False): + return self.context_ref + + +def test_build_repo_context_process_cache_refreshes_when_revision_changes(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_FILES", ["AGENTS.md"]) + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_LINES", 500) + pr_url = "https://example.com/org/repo/pull/1" + provider = RefishProvider({"AGENTS.md": "before rebase"}, pr_url=pr_url) + + first_context = build_repo_context(provider) + assert "before rebase" in first_context + + # The base branch moved (rebase/push within the TTL): the revision the cache is keyed on + # changes, so the stale entry must not be served. + provider.context_ref = "sha-2" + provider.files["AGENTS.md"] = "after rebase" + + second_context = build_repo_context(provider) + + assert "after rebase" in second_context + assert "before rebase" not in second_context + assert provider.requested_paths == ["AGENTS.md", "AGENTS.md"] + + +def test_build_repo_context_provider_cache_refreshes_when_revision_changes(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_FILES", ["AGENTS.md"]) + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_LINES", 500) + provider = RefishProvider({"AGENTS.md": "before rebase"}) + + first_context = build_repo_context(provider) + assert "before rebase" in first_context + + provider.context_ref = "sha-3" + provider.files["AGENTS.md"] = "after push" + + second_context = build_repo_context(provider) + + assert "after push" in second_context + assert provider.requested_paths == ["AGENTS.md", "AGENTS.md"] + + +def test_get_repo_context_ref_github_returns_base_sha(): + provider = GithubProvider.__new__(GithubProvider) + provider.base_url = "https://api.github.com" + provider.provider_id = None + provider.pr = Mock(base=Mock(sha="base-sha", ref="release/1.0")) + + assert provider.get_repo_context_ref() == "base-sha" + + +def test_get_repo_context_ref_github_resolves_default_branch_head(): + """Reading the default branch keys the cache on its head commit, so a push to it + invalidates cached content within the TTL instead of serving a moved commit.""" + provider = GithubProvider.__new__(GithubProvider) + provider.repo_obj = Mock() + provider.repo_obj.default_branch = "main" + provider.repo_obj.get_branch.return_value.commit.sha = "default-sha" + + assert provider.get_repo_context_ref(from_default_branch=True) == "default-sha" + provider.repo_obj.get_branch.assert_called_once_with("main") + + +def test_get_repo_context_ref_github_without_pr_base_reads_default_branch_head(): + """Without a PR base the fallback read also comes from the default branch, so the + same key rule applies.""" + provider = GithubProvider.__new__(GithubProvider) + provider.pr = None + provider.repo_obj = Mock() + provider.repo_obj.default_branch = "main" + provider.repo_obj.get_branch.return_value.commit.sha = "default-sha" + + assert provider.get_repo_context_ref() == "default-sha" + provider.repo_obj.get_branch.assert_called_once_with("main") + + +def test_get_repo_context_ref_github_falls_back_to_none_without_repo_obj(): + provider = GithubProvider.__new__(GithubProvider) + + assert provider.get_repo_context_ref(from_default_branch=True) is None + assert provider.get_repo_context_ref() is None + + +def test_build_repo_context_github_invalidates_default_branch_cache_when_head_moves(repo_context_settings): + """The shipping default reads repo-context from the default branch; a push to it + within the TTL must not serve the previous head's content. Keying the cache on the + resolved head commit closes the gap the review called out on GitHub.""" + repo_context_settings.set("CONFIG.REPO_CONTEXT_FILES", ["AGENTS.md"]) + repo_context_settings.set("CONFIG.REPO_CONTEXT_MAX_LINES", 500) + repo_context_settings.set("CONFIG.REPO_CONTEXT_FROM_DEFAULT_BRANCH", True) + provider = GithubProvider.__new__(GithubProvider) + provider.pr = Mock(base=Mock(sha="base-sha", ref="release/1.0")) + provider.repo_obj = Mock() + provider.repo_obj.default_branch = "main" + provider.repo_obj.get_contents.return_value.decoded_content = b"before push" + provider.repo_obj.get_branch.return_value.commit.sha = "head-1" + + first_context = build_repo_context(provider) + assert "before push" in first_context + + provider.repo_obj.get_contents.return_value.decoded_content = b"after push" + provider.repo_obj.get_branch.return_value.commit.sha = "head-2" + + second_context = build_repo_context(provider) + + assert "after push" in second_context + assert "before push" not in second_context + assert provider.repo_obj.get_contents.call_count == 2 + + +def test_get_repo_context_ref_gitlab_returns_target_branch(): + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "owner/repo" + provider.mr = Mock(target_branch="main") + provider.gl = Mock() + provider.gl.projects.get.return_value.default_branch = "project-default" + + assert provider.get_repo_context_ref() == "main" + assert provider.get_repo_context_ref(from_default_branch=True) == "project-default" + + +def test_get_repo_context_ref_azure_returns_base_commit(): + provider = AzureDevopsProvider.__new__(AzureDevopsProvider) + provider.workspace_slug = "my-project" + provider.repo_slug = "my-repo" + provider.pr = Mock(last_merge_target_commit=Mock(commit_id="base-sha")) + + assert provider.get_repo_context_ref() == "base-sha" + assert provider.get_repo_context_ref(from_default_branch=True) is None + + +def test_get_repo_context_ref_bitbucket_server_returns_base_commit(): + provider = BitbucketServerProvider.__new__(BitbucketServerProvider) + provider.workspace_slug = "PRJ" + provider.repo_slug = "repo" + provider.pr = SimpleNamespace(toRef={"latestCommit": "base-sha"}) + provider.bitbucket_client = Mock() + + assert provider.get_repo_context_ref() == "base-sha" + provider.bitbucket_client.get_default_branch.return_value = {"displayId": "develop"} + assert provider.get_repo_context_ref(from_default_branch=True) == "develop" + + +def test_get_repo_context_ref_gitea_returns_base_ref(): + provider = GiteaProvider.__new__(GiteaProvider) + provider.logger = Mock() + provider.owner = "owner" + provider.repo = "repo" + provider.base_sha = "base-sha" + provider.base_ref = None + provider.repo_api = Mock() + + assert provider.get_repo_context_ref() == "base-sha" + provider.repo_api.repo_get.return_value.default_branch = "main" + assert provider.get_repo_context_ref(from_default_branch=True) == "main" + + +def test_get_repo_context_ref_bitbucket_cloud_returns_destination_branch(): + provider = BitbucketProvider.__new__(BitbucketProvider) + provider.workspace_slug = "myws" + provider.repo_slug = "myrepo" + provider.pr = Mock(destination_branch="main") + provider.token = None + + assert provider.get_repo_context_ref() == "main" + provider.get_repo_default_branch = Mock(return_value="develop") + assert provider.get_repo_context_ref(from_default_branch=True) == "develop" + + +def test_get_repo_context_ref_inherits_none_for_providers_without_repo_context(): + provider = CodeCommitProvider.__new__(CodeCommitProvider) + + assert provider.get_repo_context_ref() is None + assert provider.get_repo_context_ref(from_default_branch=True) is None + + +@pytest.mark.parametrize("provider_class", [GithubProvider, GitLabProvider]) +@pytest.mark.parametrize("allowlist", [[], "group/lib", None]) +def test_sibling_reads_require_host_approval(repo_context_settings, provider_class, allowlist): + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", allowlist) + provider = provider_class.__new__(provider_class) + provider.github_client = Mock() + provider.gl = Mock() + + assert provider.get_sibling_repo_file_content("group/lib", "api.py") == "" + provider.github_client.get_repo.assert_not_called() + provider.gl.projects.get.assert_not_called() + + +@pytest.mark.parametrize("current_id", ["123", "group/app/current"]) +@pytest.mark.parametrize("sibling_id", ["456", "group/libraries/lib"]) +def test_gitlab_siblings_resolve_numeric_ids_and_subgroups(repo_context_settings, current_id, sibling_id): + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", [sibling_id]) + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = current_id + current = SimpleNamespace(path_with_namespace="group/app/current") + sibling = Mock(id=456, path_with_namespace="group/libraries/lib", visibility="public", default_branch="main") + sibling.files.get.return_value.decode.return_value = b"interface" + provider.gl = Mock() + provider.gl.projects.get.side_effect = lambda identifier: current if identifier == current_id else sibling + + assert provider.get_sibling_repo_file_content(sibling_id, "api.py") == "interface" + sibling.files.get.assert_called_once_with(file_path="api.py", ref="main") + + +@pytest.mark.parametrize("sibling_id,resolved_path", [ + ("group/lib", "other/lib"), + ("group/lib", "group/renamed"), + ("456", "other/lib"), +]) +def test_gitlab_rejects_resolved_identity_or_namespace_change(repo_context_settings, sibling_id, resolved_path): + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", [sibling_id]) + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/current" + current = SimpleNamespace(path_with_namespace="group/current") + sibling = Mock(id=456, path_with_namespace=resolved_path, visibility="public") + provider.gl = Mock() + provider.gl.projects.get.side_effect = lambda identifier: current if identifier == provider.id_project else sibling + + assert provider.get_sibling_repo_file_content(sibling_id, "api.py") == "" + sibling.files.get.assert_not_called() + + +@pytest.mark.parametrize("current_path,sibling_path,expected", [ + ("MyOrg/current", "myorg/lib", "interface"), + ("other/current", "myorg/lib", ""), + ("myorg/current", "other/lib", ""), + ("myorg/current", "myorg/renamed", ""), +]) +def test_github_checks_resolved_owners_and_identity(repo_context_settings, current_path, sibling_path, expected): + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", ["MyOrg/lib"]) + provider = GithubProvider.__new__(GithubProvider) + provider.repo = "myorg/current" + current = SimpleNamespace(full_name=current_path) + sibling = Mock(full_name=sibling_path, private=False, visibility="public") + sibling.get_contents.return_value.decoded_content = b"interface" + provider.github_client = Mock() + provider.github_client.get_repo.side_effect = lambda identifier: current if identifier == provider.repo else sibling + + assert provider.get_sibling_repo_file_content("myorg/lib", "api.py") == expected + if not expected: + sibling.get_contents.assert_not_called() + + +def test_selected_sibling_build_requires_allowlist(repo_context_settings): + repo_context_settings.set("CONFIG.REPO_CONTEXT_FILES", [ + {"repo_id": "group/lib", "file_path": "api.py"}, + ]) + provider = GitLabProvider.__new__(GitLabProvider) + provider.id_project = "group/current" + provider.gl = Mock() + sibling = Mock(id=456, path_with_namespace="group/lib", visibility="public", default_branch="main") + sibling.files.get.return_value.decode.return_value = b"interface" + provider.gl.projects.get.side_effect = lambda identifier: ( + SimpleNamespace(path_with_namespace="group/current") if identifier == provider.id_project else sibling + ) + + assert build_repo_context(provider) == "" + provider.gl.projects.get.assert_not_called() + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", ["group/lib"]) + assert "interface" in build_repo_context(provider) + repo_context_settings.set("CONFIG.REPO_CONTEXT_SIBLING_REPOS", []) + assert build_repo_context(provider) == "" diff --git a/tests/unittest/test_review_comment_identity.py b/tests/unittest/test_review_comment_identity.py index edb9561280..8627e3a9f6 100644 --- a/tests/unittest/test_review_comment_identity.py +++ b/tests/unittest/test_review_comment_identity.py @@ -13,6 +13,7 @@ ) from pr_agent.config_loader import get_settings from pr_agent.git_providers.azuredevops_provider import AzureDevopsProvider +from pr_agent.git_providers.bitbucket_provider import BitbucketProvider from pr_agent.git_providers.gitea_provider import GiteaProvider from pr_agent.git_providers.github_provider import GithubProvider from tests.unittest._settings_helpers import restore_settings, snapshot_settings @@ -174,11 +175,35 @@ def test_github_comment_path_forwards_review_identity(): restore_settings(snapshot) assert provider.publish_persistent_comment_full.call_args.kwargs == { + "as_thread": False, "identity_marker": PRReviewIdentity.REGULAR.value, "legacy_initial_header": legacy_header, } +def test_github_comment_path_forwards_as_thread(): + snapshot = snapshot_settings(["github.publish_as_check_run"]) + provider = GithubProvider.__new__(GithubProvider) + provider.publish_persistent_comment_full = MagicMock() + review = "## Team Review 🔍\n\nbody" + try: + get_settings().set("github.publish_as_check_run", False) + + provider.publish_persistent_comment( + review, + initial_header="## Team Review 🔍", + as_thread=True, + ) + finally: + restore_settings(snapshot) + + assert provider.publish_persistent_comment_full.call_args.kwargs == { + "as_thread": True, + "identity_marker": None, + "legacy_initial_header": None, + } + + def test_azure_comment_path_forwards_review_identity(): provider = AzureDevopsProvider.__new__(AzureDevopsProvider) provider.publish_persistent_comment_full = MagicMock() @@ -193,11 +218,33 @@ def test_azure_comment_path_forwards_review_identity(): ) assert provider.publish_persistent_comment_full.call_args.kwargs == { + "as_thread": False, "identity_marker": PRReviewIdentity.REGULAR.value, "legacy_initial_header": legacy_header, } +def test_bitbucket_comment_path_forwards_review_identity(): + provider = BitbucketProvider.__new__(BitbucketProvider) + provider.pr = MagicMock() + provider.pr.comments.return_value = [] + provider.publish_comment = MagicMock() + review = "## Team Review 🔍\n\nbody" + legacy_header = "## PR Reviewer Guide 🔍" + + provider.publish_persistent_comment( + review, + initial_header="## Team Review 🔍", + identity_marker=PRReviewIdentity.REGULAR.value, + legacy_initial_header=legacy_header, + ) + + assert provider.supports_review_comment_identity() is True + provider.publish_comment.assert_called_once_with( + "## Team Review 🔍\n\n\n\nbody" + ) + + def test_gitea_keeps_identity_inactive_but_preserves_wrapper_arguments(): provider = GiteaProvider.__new__(GiteaProvider) published = object() @@ -215,6 +262,7 @@ def test_gitea_keeps_identity_inactive_but_preserves_wrapper_arguments(): assert provider.supports_review_comment_identity() is False assert result is published assert provider.publish_persistent_comment_full.call_args.kwargs == { + "as_thread": False, "identity_marker": PRReviewIdentity.REGULAR.value, "legacy_initial_header": legacy_header, } diff --git a/tests/unittest/test_review_finding_persistence.py b/tests/unittest/test_review_finding_persistence.py index 25c3acac4f..10e841527a 100644 --- a/tests/unittest/test_review_finding_persistence.py +++ b/tests/unittest/test_review_finding_persistence.py @@ -24,6 +24,37 @@ def _reviewer(provider): return reviewer +class _PersistentCommentProvider: + publish_persistent_comment = GitProvider.publish_persistent_comment + publish_persistent_comment_full = GitProvider.publish_persistent_comment_full + supports_comment_editing = GitProvider.supports_comment_editing + + def __init__(self, comments): + self.comments = comments + self.edited = [] + self.published = [] + + def get_issue_comments(self): + return self.comments + + def get_issue_comments_newest_first(self): + return list(reversed(self.comments)) + + def get_latest_commit_url(self): + return "commit-url" + + def get_comment_url(self, comment): + return "comment-url" + + def edit_comment(self, comment, body): + self.edited.append((comment, body)) + return True + + def publish_comment(self, body, **kwargs): + self.published.append((body, kwargs)) + return "published" + + def test_invalid_structured_finding_fails_closed(monkeypatch): settings = get_settings() monkeypatch.setattr(settings.config, "publish_output", True) @@ -150,7 +181,7 @@ def test_stateful_persistent_update_still_creates_first_comment(): provider.publish_comment.assert_called_once_with("new review") -def test_stateful_mode_is_disabled_for_generic_persistent_publisher(monkeypatch): +def test_stateful_mode_is_enabled_for_generic_persistent_publisher(monkeypatch): settings = get_settings() monkeypatch.setattr(settings.config, "publish_output", True) monkeypatch.setattr(settings.pr_reviewer, "persistent_comment", True) @@ -159,7 +190,36 @@ def test_stateful_mode_is_disabled_for_generic_persistent_publisher(monkeypatch) provider.publish_persistent_comment = GitProvider.publish_persistent_comment.__get__(provider, type(provider)) provider.is_supported.return_value = True reviewer = _reviewer(provider) - assert reviewer._review_finding_state_enabled() is False + assert reviewer._review_finding_state_enabled() is True + + +def test_base_persistent_comment_updates_matching_comment(): + header = "## PR Reviewer Guide 🔍" + comment = SimpleNamespace(body=f"{header}\n\nold review") + provider = _PersistentCommentProvider([comment]) + + result = provider.publish_persistent_comment( + f"{header}\n\nnew review", + initial_header=header, + update_header=False, + final_update_message=False, + ) + + assert result is comment + assert provider.edited == [(comment, f"{header}\n\nnew review")] + assert provider.published == [] + + +def test_base_persistent_comment_creates_comment_when_no_match_exists(): + provider = _PersistentCommentProvider([]) + + result = provider.publish_persistent_comment( + "new review", + initial_header="## PR Reviewer Guide 🔍", + ) + + assert result == "published" + assert provider.published == [("new review", {})] def test_malformed_state_marker_is_replaced_without_duplicate_comment(): @@ -249,3 +309,40 @@ def test_dict_comment_edit_failure_does_not_fallback(): assert result is None provider.edit_comment.assert_called_once_with(comment, "new review") provider.publish_comment.assert_not_called() + + +class _NoEditProvider(_PersistentCommentProvider): + """A provider that never implemented edit_comment, like gerrit and codecommit. + + The inherited base edit_comment is a no-op returning None, which + publish_persistent_comment_full cannot tell apart from a successful edit. + """ + + edit_comment = GitProvider.edit_comment + supports_comment_editing = GitProvider.supports_comment_editing + + +def test_provider_without_edit_comment_publishes_the_body_instead_of_losing_it(): + header = "## PR Reviewer Guide 🔍" + comment = SimpleNamespace(body=f"{header}\n\nold review") + provider = _NoEditProvider([comment]) + + assert provider.supports_comment_editing() is False + + result = provider.publish_persistent_comment( + f"{header}\n\nnew review", + initial_header=header, + update_header=False, + final_update_message=False, + ) + + # The review body itself must reach the PR. Delegating to the full + # implementation here would no-op the edit and publish only the + # "updated to latest commit" status line, silently dropping the review. + assert result == "published" + assert provider.published == [(f"{header}\n\nnew review", {})] + + +def test_provider_with_edit_comment_reports_editing_support(): + provider = _PersistentCommentProvider([]) + assert GitProvider.supports_comment_editing(provider) is True diff --git a/tests/unittest/test_review_large_diff_chunking.py b/tests/unittest/test_review_large_diff_chunking.py index 87b36d020a..dae7ee2008 100644 --- a/tests/unittest/test_review_large_diff_chunking.py +++ b/tests/unittest/test_review_large_diff_chunking.py @@ -10,6 +10,8 @@ import pytest +from pr_agent.algo.pr_processing import PreparedPRDiff +from pr_agent.algo.review_finding_state import ParsedReviewState, reconcile_review_findings from pr_agent.config_loader import get_settings from pr_agent.tools.pr_reviewer import PRReviewer from tests.unittest._settings_helpers import restore_settings, snapshot_settings @@ -123,6 +125,29 @@ async def test_a_truncated_diff_is_reviewed_chunk_by_chunk_and_merged(chunking_e assert reviewer.remaining_files_list == ["still_left_out.py"] +@pytest.mark.asyncio +async def test_chunked_review_reuses_the_prepared_diff_for_the_same_model_attempt(chunking_enabled): + reviewer = _make_reviewer() + reviewer._get_prediction = AsyncMock(side_effect=[CHUNK_A, CHUNK_B]) + prepared = PreparedPRDiff( + diff="first compressed diff", + remaining_files_list=["b.py"], + file_dict={"a.py": {"patch": "chunk-a", "tokens": 10}}, + ) + + with ( + patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value=prepared) as get_pr_diff, + patch("pr_agent.tools.pr_reviewer.get_pr_multi_diffs", + return_value=(["chunk-a", "chunk-b"], ["still_left_out.py"])) as get_pr_multi_diffs, + ): + await reviewer._prepare_prediction("model") + + assert get_pr_diff.call_args.kwargs["return_prepared"] is True + assert get_pr_multi_diffs.call_args.kwargs["prepared_diff"] is prepared + assert reviewer.review_chunk_count == 2 + assert reviewer.remaining_files_list == ["still_left_out.py"] + + @pytest.mark.asyncio async def test_max_number_of_calls_bounds_the_number_of_chunks(chunking_enabled): get_settings().set("pr_reviewer.max_number_of_calls", 7) @@ -166,16 +191,56 @@ async def test_a_chunk_that_fails_does_not_lose_the_chunks_that_succeeded(chunki patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value=("diff", ["b.py"])), patch("pr_agent.tools.pr_reviewer.get_pr_multi_diffs", return_value=(["chunk-a", "chunk-b"], [])), + pytest.raises(RuntimeError, match="model refused"), ): await reviewer._prepare_prediction("model") + reviewer._get_prediction.side_effect = [CHUNK_A] + await reviewer._prepare_prediction("fallback-model") + assert reviewer.prediction_data["review"]["score"] == "40" assert reviewer.review_chunk_count == 2 - assert reviewer.review_failed_chunk_count == 1 + assert reviewer.review_failed_chunk_count == 0 + assert [call.args[1] for call in reviewer._get_prediction.await_args_list] == ["chunk-a", "chunk-b", "chunk-a"] + + +@pytest.mark.asyncio +async def test_a_failed_chunk_blocks_persistent_finding_resolution(chunking_enabled): + """A partial review must stay partial for the finding-state lifecycle too.""" + reviewer = _make_reviewer() + reviewer._get_prediction = AsyncMock(side_effect=[RuntimeError("model refused"), CHUNK_B]) + + with ( + patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value=("diff", ["b.py"])), + patch("pr_agent.tools.pr_reviewer.get_pr_multi_diffs", + return_value=(["chunk-a", "chunk-b"], [])), + ): + await reviewer._prepare_prediction("model") + + previous_state = reconcile_review_findings( + None, + [{"path": "a.py", "body": "old finding", "line_start": 3, "line_end": 4}], + allow_resolution=False, + head_sha="head-1", + timestamp="2026-09-09T00:00:00+00:00", + ).state + reviewer._review_finding_state_enabled = lambda: True + reviewer._load_review_finding_state = lambda: ParsedReviewState( + previous_state, present=True, valid=True + ) + reviewer._review_head_sha = lambda: "head-2" + reviewer._review_run_id = lambda: "run-2" + + reviewer._prepare_review_finding_state(reviewer.prediction_data) + + assert reviewer._review_state_result is not None + assert reviewer._review_state_result.resolved_ids == () + assert reviewer._review_state_result.state["last_run"]["complete"] is False + assert reviewer._review_state_result.state["findings"][0]["state"] == "ACTIVE" @pytest.mark.asyncio -async def test_an_empty_chunk_does_not_lose_a_valid_sibling_or_trigger_fallback(chunking_enabled): +async def test_a_malformed_chunk_is_retried_without_repeating_successful_chunks(chunking_enabled): reviewer = _make_reviewer() reviewer._get_prediction = AsyncMock(side_effect=["review: {}", CHUNK_B]) @@ -183,13 +248,18 @@ async def test_an_empty_chunk_does_not_lose_a_valid_sibling_or_trigger_fallback( patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value=("diff", ["b.py"])), patch("pr_agent.tools.pr_reviewer.get_pr_multi_diffs", return_value=(["chunk-a", "chunk-b"], [])), + pytest.raises(ValueError, match="non-empty review"), ): await reviewer._prepare_prediction("model") assert reviewer._get_prediction.await_count == 2 + assert reviewer.prediction_data is None + + reviewer._get_prediction.side_effect = [CHUNK_A] + await reviewer._prepare_prediction("fallback-model") + assert reviewer.prediction_data["review"]["score"] == "40" - assert reviewer.review_chunk_count == 2 - assert reviewer.review_failed_chunk_count == 1 + assert [call.args[1] for call in reviewer._get_prediction.await_args_list] == ["chunk-a", "chunk-b", "chunk-a"] @pytest.mark.asyncio @@ -212,20 +282,20 @@ async def test_a_review_where_every_chunk_failed_raises_so_a_fallback_model_is_t ["not yaml at all", "nor is this"], ["review: {}", "review: {}"], ]) -async def test_chunks_without_nonempty_reviews_fall_back_to_a_single_call_review(chunking_enabled, - chunk_predictions): +async def test_chunks_without_nonempty_reviews_fail_the_model_attempt(chunking_enabled, + chunk_predictions): reviewer = _make_reviewer() - reviewer._get_prediction = AsyncMock(side_effect=[*chunk_predictions, CHUNK_A]) + reviewer._get_prediction = AsyncMock(side_effect=chunk_predictions) with ( patch("pr_agent.tools.pr_reviewer.get_pr_diff", return_value=("diff", ["b.py"])), patch("pr_agent.tools.pr_reviewer.get_pr_multi_diffs", return_value=(["chunk-a", "chunk-b"], [])), + pytest.raises(ValueError, match="non-empty review"), ): await reviewer._prepare_prediction("model") - assert reviewer._get_prediction.await_count == 3 - assert reviewer.prediction == CHUNK_A + assert reviewer._get_prediction.await_count == 2 assert reviewer.prediction_data is None diff --git a/tests/unittest/test_review_persistent_comment_authorship.py b/tests/unittest/test_review_persistent_comment_authorship.py index f457ab2467..00a073df1d 100644 --- a/tests/unittest/test_review_persistent_comment_authorship.py +++ b/tests/unittest/test_review_persistent_comment_authorship.py @@ -17,7 +17,7 @@ from unittest.mock import MagicMock import pytest -from github import GithubException +from github import Auth, GithubException from pr_agent.algo.review_finding_state import append_review_state, reconcile_review_findings from pr_agent.algo.types import EDIT_TYPE, FilePatchInfo @@ -367,9 +367,11 @@ def test_a_transient_app_login_failure_is_retried(monkeypatch): monkeypatch.setattr(get_settings(), "github", SimpleNamespace( app_id="1", private_key="key", deployment_type="app"), raising=False) attempts = [] + auths = [] def integration(**kwargs): attempts.append(1) + auths.append(kwargs["auth"]) if len(attempts) == 1: raise RuntimeError("connection reset") return SimpleNamespace(get_app=lambda: SimpleNamespace(slug="pr-agent")) @@ -379,6 +381,7 @@ def integration(**kwargs): assert provider._resolve_app_login() == "" assert provider._resolve_app_login() == "pr-agent[bot]" assert len(attempts) == 2 + assert all(isinstance(auth, Auth.AppAuth) for auth in auths) def test_a_resolved_app_login_is_cached(monkeypatch): @@ -389,9 +392,11 @@ def test_a_resolved_app_login_is_cached(monkeypatch): monkeypatch.setattr(get_settings(), "github", SimpleNamespace( app_id="1", private_key="key", deployment_type="app"), raising=False) attempts = [] + auths = [] def integration(**kwargs): attempts.append(1) + auths.append(kwargs["auth"]) return SimpleNamespace(get_app=lambda: SimpleNamespace(slug="pr-agent")) monkeypatch.setattr("pr_agent.git_providers.github_provider.GithubIntegration", integration) @@ -399,6 +404,7 @@ def integration(**kwargs): assert provider._resolve_app_login() == "pr-agent[bot]" assert provider._resolve_app_login() == "pr-agent[bot]" assert len(attempts) == 1 + assert all(isinstance(auth, Auth.AppAuth) for auth in auths) def test_an_app_that_never_resolves_stays_unproven(monkeypatch): diff --git a/tests/unittest/test_run_details_wiring.py b/tests/unittest/test_run_details_wiring.py index b9b5684830..5f47fc0c26 100644 --- a/tests/unittest/test_run_details_wiring.py +++ b/tests/unittest/test_run_details_wiring.py @@ -4,6 +4,7 @@ import pytest from pr_agent.algo.run_details import init_run_details, record_ai_call, record_model_used +from pr_agent.algo.types import FilePatchInfo from pr_agent.config_loader import get_settings from pr_agent.tools.pr_code_suggestions import PRCodeSuggestions from pr_agent.tools.pr_description import PRDescription @@ -137,7 +138,7 @@ async def test_pr_description_appends_run_details_only_when_enabled(monkeypatch) description.git_provider.is_supported.side_effect = lambda cap: cap == "gfm_markdown" description._prepare_data = MagicMock() - description._prepare_pr_answer = MagicMock(return_value=("AI title", "Base description body", "", [])) + description._prepare_pr_answer = MagicMock(return_value=("AI title", "Base description body", "")) monkeypatch.setattr("pr_agent.tools.pr_description.init_run_details", _seeded_init_run_details) monkeypatch.setattr("pr_agent.tools.pr_description.extract_and_cache_pr_tickets", _noop_async) @@ -176,10 +177,21 @@ async def test_pr_code_suggestions_appends_run_details_only_when_enabled(monkeyp suggestions.git_provider = MagicMock() suggestions.git_provider.get_files.return_value = ["changed.py"] suggestions.git_provider.is_supported.side_effect = lambda cap: cap == "gfm_markdown" + suggestions.git_provider.diff_files = [ + FilePatchInfo( + base_file="old()\n", + head_file="new()\n", + patch="@@ -1,1 +1,1 @@\n-old()\n+new()\n", + filename="changed.py", + ) + ] suggestions.generate_summarized_suggestions = MagicMock(return_value="Base suggestions body") async def _fake_retry(*_args, **_kwargs): - return {"code_suggestions": [{"label": "style"}]} + return {"code_suggestions": [ + {"label": "style", "relevant_file": "changed.py", + "relevant_lines_start": 1, "relevant_lines_end": 1}, + ]} monkeypatch.setattr("pr_agent.tools.pr_code_suggestions.init_run_details", _seeded_init_run_details) monkeypatch.setattr("pr_agent.tools.pr_code_suggestions.retry_with_fallback_models", _fake_retry) diff --git a/tests/unittest/test_similar_issue_embeddings.py b/tests/unittest/test_similar_issue_embeddings.py index a09e880a24..c7035e6df7 100644 --- a/tests/unittest/test_similar_issue_embeddings.py +++ b/tests/unittest/test_similar_issue_embeddings.py @@ -1,6 +1,5 @@ """Embed with the openai>=1.0 client that requirements.txt pins.""" import inspect -import sys from types import SimpleNamespace import pytest @@ -102,7 +101,7 @@ def __init__(self, api_key=None): @pytest.mark.asyncio -async def test_pinecone_query_import_is_available_in_run(monkeypatch): +async def test_pinecone_query_branch_runs(monkeypatch): queried = [] class FakeIndex: @@ -110,7 +109,6 @@ def query(self, *args, **kwargs): queried.append(kwargs) return SimpleNamespace(to_dict=lambda: {"matches": []}) - monkeypatch.setitem(sys.modules, "pinecone", SimpleNamespace(Index=lambda **kwargs: FakeIndex())) monkeypatch.setattr(psi, "_embed", lambda texts: [[0.5]]) monkeypatch.setattr( psi, @@ -127,6 +125,7 @@ def query(self, *args, **kwargs): tool.issue_url = "https://github.com/example/repo/issues/1" tool.index_name = "issues" tool.repo_name_for_index = "example-repo" + tool.pc = SimpleNamespace(Index=lambda name: FakeIndex()) tool.git_provider = SimpleNamespace( _parse_issue_url=lambda url: ("example/repo", 1), repo_obj=SimpleNamespace(get_issue=lambda number: issue), @@ -134,3 +133,4 @@ def query(self, *args, **kwargs): assert await tool.run() is None assert queried, "run() never entered the pinecone branch" + assert queried[0]["vector"] == [0.5] diff --git a/tests/unittest/test_similar_issue_non_github.py b/tests/unittest/test_similar_issue_non_github.py deleted file mode 100644 index 6581448692..0000000000 --- a/tests/unittest/test_similar_issue_non_github.py +++ /dev/null @@ -1,49 +0,0 @@ -import pytest - -from pr_agent.tools.pr_similar_issue import PRSimilarIssue - - -@pytest.mark.asyncio -async def test_similar_issue_non_github_publishes_message(monkeypatch): - class FakeProvider: - def __init__(self): - self.comments = [] - - def publish_comment(self, body): - self.comments.append(body) - - fake_provider = FakeProvider() - - class FakeSettings: - class config: - git_provider = "gitlab" - publish_output = True - - monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_settings", lambda: FakeSettings) - monkeypatch.setattr( - "pr_agent.git_providers.get_git_provider_with_context", - lambda _: fake_provider, - ) - - tool = PRSimilarIssue("https://gitlab.example.com/group/repo/-/merge_requests/1", None) - result = await tool.run() - - assert result == "" - assert fake_provider.comments == [ - "The /similar_issue tool is currently supported only for GitHub." - ] - - -@pytest.mark.asyncio -async def test_similar_issue_non_github_no_publish(monkeypatch): - class FakeSettings: - class config: - git_provider = "gitlab" - publish_output = False - - monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_settings", lambda: FakeSettings) - - tool = PRSimilarIssue("https://gitlab.example.com/group/repo/-/merge_requests/1", None) - result = await tool.run() - - assert result == "" diff --git a/tests/unittest/test_similar_issue_pinecone.py b/tests/unittest/test_similar_issue_pinecone.py new file mode 100644 index 0000000000..4456afa1c3 --- /dev/null +++ b/tests/unittest/test_similar_issue_pinecone.py @@ -0,0 +1,150 @@ +"""Focused unit tests for the pinecone backend of the similar issue tool.""" +import sys +from pathlib import Path +from types import SimpleNamespace + +import pr_agent.tools.pr_similar_issue as psi + + +class SettingsStub: + class pr_similar_issue: + skip_comments = True + + +class _PandasSeries(list): + @property + def values(self): + return list(self) + + def to_list(self): + return list(self) + + +class _PandasDataFrame: + def __init__(self, records): + self._records = list(records) + self._overrides = {} + + def __getitem__(self, column): + if column in self._overrides: + return _PandasSeries(self._overrides[column]) + return _PandasSeries([row[column] for row in self._records]) + + def __setitem__(self, column, values): + self._overrides[column] = list(values) + + def to_dict(self, orient="records"): + rows = [dict(record) for record in self._records] + for column, values in self._overrides.items(): + for row, value in zip(rows, values, strict=True): + row[column] = value + return rows + + +def _make_issue(number): + return SimpleNamespace( + pull_request=False, + title=f"Issue {number}", + body="Body", + number=number, + user=SimpleNamespace(login="user"), + created_at="2020-01-01", + get_comments=lambda: [], + ) + + +def _make_tool(pc): + tool = psi.PRSimilarIssue.__new__(psi.PRSimilarIssue) + tool.pc = pc + tool.pc_spec = None + tool.index_name = "codium-ai-pr-agent-issues" + tool.max_issues_to_scan = 100 + tool.token_handler = SimpleNamespace(count_tokens=lambda text: 1) + tool.pinecone_index = None + return tool + + +def _stub_embeddings(monkeypatch): + monkeypatch.setitem(sys.modules, "pandas", + SimpleNamespace(DataFrame=_PandasDataFrame)) + embedding = [0.5, 0.5] + monkeypatch.setattr(psi, "get_settings", lambda: SettingsStub) + monkeypatch.setattr(psi, "get_max_tokens", lambda model: 8192) + monkeypatch.setattr(psi, "_embed_with_fallback", + lambda texts: [embedding for _ in texts]) + monkeypatch.setattr(psi.time, "sleep", lambda seconds: None) + + +def test_pinecone_upsert_passes_batch_size_and_max_concurrency(monkeypatch): + calls = {} + + class FakeIndex: + def upsert(self, **kwargs): + calls.update(kwargs) + + tool = _make_tool(SimpleNamespace(Index=lambda name: FakeIndex())) + _stub_embeddings(monkeypatch) + + tool._update_index_with_issues([_make_issue(7)], "example-repo", upsert=True) + + assert calls["namespace"] == "" + assert calls["batch_size"] == 100 + assert calls["max_concurrency"] == 10 + assert [vector[0] for vector in calls["vectors"]] == [ + "example_issue_example-repo", + "issue_7.issue", + ] + + +def test_pinecone_upsert_path_skips_index_creation(monkeypatch): + created = [] + + class FakeIndex: + def upsert(self, **kwargs): + pass + + pc = SimpleNamespace( + Index=lambda name: FakeIndex(), + create_index=lambda **kwargs: created.append(kwargs), + ) + tool = _make_tool(pc) + _stub_embeddings(monkeypatch) + + tool._update_index_with_issues([_make_issue(7)], "example-repo", upsert=True) + + assert created == [] + + +def test_pinecone_create_index_path_builds_new_index_then_upserts(monkeypatch): + created = [] + upserted = [] + + class FakeIndex: + def upsert(self, **kwargs): + upserted.append(kwargs) + + pc = SimpleNamespace( + Index=lambda name: FakeIndex(), + create_index=lambda **kwargs: created.append(kwargs), + ) + tool = _make_tool(pc) + _stub_embeddings(monkeypatch) + + tool._update_index_with_issues([_make_issue(7)], "example-repo", upsert=False) + + assert len(created) == 1 + kwargs = created[0] + assert kwargs["name"] == tool.index_name + assert kwargs["dimension"] == 2 + assert kwargs["metric"] == "cosine" + assert kwargs["timeout"] == 120 + assert upserted, "expected the upsert to run after index creation" + + +def test_vectordb_defaults_to_lancedb(): + config_path = Path(psi.__file__).resolve().parents[2] / "pr_agent" / "settings" / "configuration.toml" + content = config_path.read_text(encoding="utf-8") + section = content.split("[pr_similar_issue]")[1].split("[")[0] + + assert 'vectordb = "lancedb"' in section + assert 'vectordb = "pinecone"' not in section diff --git a/tests/unittest/test_similar_issue_provider_capability.py b/tests/unittest/test_similar_issue_provider_capability.py new file mode 100644 index 0000000000..db50b52aa1 --- /dev/null +++ b/tests/unittest/test_similar_issue_provider_capability.py @@ -0,0 +1,134 @@ +from types import SimpleNamespace + +import pytest + +from pr_agent.git_providers.git_provider import GitProvider +from pr_agent.git_providers.github_provider import GithubProvider +from pr_agent.tools.pr_similar_issue import PRSimilarIssue, _provider_supports_issue_indexing + + +class UnsupportedProvider(GitProvider): + """A provider class that does not declare the issue-indexing capability.""" + + def is_supported(self, capability: str) -> bool: + return False + + +@pytest.mark.asyncio +async def test_similar_issue_unsupported_provider_publishes_message(monkeypatch): + class FakeProvider: + def __init__(self): + self.comments = [] + + def publish_comment(self, body): + self.comments.append(body) + + fake_provider = FakeProvider() + + class FakeSettings: + class config: + git_provider = "gitlab" + publish_output = True + + monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_settings", lambda: FakeSettings) + monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_git_provider", lambda: UnsupportedProvider) + monkeypatch.setattr( + "pr_agent.git_providers.get_git_provider_with_context", + lambda _: fake_provider, + ) + + tool = PRSimilarIssue("https://gitlab.example.com/group/repo/-/merge_requests/1", None) + result = await tool.run() + + assert result == "" + assert fake_provider.comments == [ + "The /similar_issue tool is not supported by the configured git provider." + ] + + +@pytest.mark.asyncio +async def test_similar_issue_unsupported_provider_no_publish(monkeypatch): + class FakeSettings: + class config: + git_provider = "gitlab" + publish_output = False + + monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_settings", lambda: FakeSettings) + monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_git_provider", lambda: UnsupportedProvider) + + tool = PRSimilarIssue("https://gitlab.example.com/group/repo/-/merge_requests/1", None) + result = await tool.run() + + assert result == "" + + +def test_github_provider_declares_issue_indexing(): + """The capability replaces the previous `git_provider == "github"` string comparison.""" + assert GithubProvider.supports_issue_indexing() is True + + +def test_git_provider_default_is_unsupported(): + """A provider that does not override the capability is treated as unsupported.""" + assert GitProvider.supports_issue_indexing() is False + assert UnsupportedProvider.supports_issue_indexing() is False + + +def test_capability_is_read_off_the_registered_class(monkeypatch): + """A provider registered from outside the package is judged by what it declares.""" + + class ThirdPartyProvider(GitProvider): + def is_supported(self, capability: str) -> bool: + return False + + @classmethod + def supports_issue_indexing(cls) -> bool: + return True + + monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_git_provider", lambda: ThirdPartyProvider) + assert _provider_supports_issue_indexing() is True + + +def test_declared_capability_beats_the_configured_provider_id(monkeypatch): + """A non-`github` id that declares the capability is supported. + + This is the behavior change: the previous `config.git_provider == "github"` comparison + turned any other id away regardless of what the provider could do. The provider is built and + interrupted at its GitHub client hook, so the test observes that the guard let it through + without running the indexing path. + """ + + class Sentinel(Exception): + pass + + class CapableProvider(GithubProvider): + def _get_github_client(self): + raise Sentinel + + @classmethod + def supports_issue_indexing(cls) -> bool: + return True + + class FakeSettings: + class config: + git_provider = "my-forge" + publish_output = False + CLI_MODE = True + + CONFIG = config + pr_similar_issue = SimpleNamespace(max_issues_to_scan=10, vectordb="pinecone") + + monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_settings", lambda: FakeSettings) + monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_git_provider", lambda: CapableProvider) + + with pytest.raises(Sentinel): + PRSimilarIssue("https://my-forge.example.com/group/repo/issues/1", None) + + +def test_unknown_provider_configuration_is_unsupported(monkeypatch): + """An unresolvable provider id is reported as unsupported rather than raised.""" + + def raise_unknown(): + raise ValueError("Unknown git provider: nope") + + monkeypatch.setattr("pr_agent.tools.pr_similar_issue.get_git_provider", raise_unknown) + assert _provider_supports_issue_indexing() is False diff --git a/tests/unittest/test_similar_issue_qdrant_collection.py b/tests/unittest/test_similar_issue_qdrant_collection.py index 8878a78f36..977f5645fe 100644 --- a/tests/unittest/test_similar_issue_qdrant_collection.py +++ b/tests/unittest/test_similar_issue_qdrant_collection.py @@ -23,7 +23,7 @@ def test_suffix_is_scoped_to_qdrant_only(): assert call_site in source assert 'index_name = self.index_name = "codium-ai-pr-agent-issues"' in source - assert "pinecone.Index(index_name=self.index_name)" in source + assert "self.pc.Index(name=self.index_name)" in source assert "self.db.create_table(self.index_name, data=df, mode=\"overwrite\")" in source assert "self.qdrant_collection_name" not in source.split("elif get_settings().pr_similar_issue.vectordb == \"qdrant\":")[0] diff --git a/tests/unittest/test_suggestion_chunk_recovery.py b/tests/unittest/test_suggestion_chunk_recovery.py new file mode 100644 index 0000000000..b6d4f44a80 --- /dev/null +++ b/tests/unittest/test_suggestion_chunk_recovery.py @@ -0,0 +1,480 @@ +"""Behavior of recovery after a partially successful suggestion pass.""" + +import asyncio +import copy +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import yaml +from starlette_context import request_cycle_context + +import pr_agent.tools.pr_code_suggestions as module +from pr_agent.algo.pr_processing import retry_with_fallback_models +from pr_agent.algo.run_details import get_run_details, init_run_details +from pr_agent.algo.types import FilePatchInfo +from pr_agent.config_loader import get_settings +from tests.unittest._settings_helpers import restore_settings, snapshot_settings + +PRCodeSuggestions = module.PRCodeSuggestions + + +@pytest.fixture +def configured(): + values = { + "config.model": "gpt-4o", + "config.fallback_models": ["gpt-4o-mini", "gpt-4.1"], + "config.publish_output": True, + "openai.deployment_id": "primary", + "openai.fallback_deployments": ["secondary", "last"], + "pr_code_suggestions.decouple_hunks": True, + "pr_code_suggestions.parallel_calls": True, + "pr_code_suggestions.suggestions_score_threshold": 0, + "pr_code_suggestions.max_suggestions_per_file": 0, + # Not used by most tests here, but snapshotted so the routed-primary test's + # writes are restored and no routing state leaks into later tests. + "model_routing.enable": False, + "model_routing.rules": [], + } + snapshot = snapshot_settings(tuple(values)) + for key, value in values.items(): + get_settings().set(key, value) + yield + restore_settings(snapshot) + + +def make_tool(monkeypatch, failures): + tool = PRCodeSuggestions.__new__(PRCodeSuggestions) + tool.git_provider = MagicMock() + tool.token_handler = MagicMock() + tool.vars = {"instructions": "scope-one"} + tool.pr_code_suggestions_prompt_system = "Review {{ instructions }}" + tool.pr_code_suggestions_prompt_user = "{{ diff }}" + calls = [] + + async def completion(*, model, system, user, **kwargs): + calls.append((model, user, get_settings().get("openai.deployment_id"), system)) + failure = failures.get((model, user)) + if isinstance(failure, BaseException): + raise failure + if callable(failure): + await failure() + if isinstance(failure, str): + return failure, "stop" + suggestion = { + "relevant_file": user + ".py", "one_sentence_summary": "Check " + user, + "label": "possible issue", "suggestion_content": "Handle the error.", + "existing_code": "old()", "improved_code": "new()", + "relevant_lines_start": 1, "relevant_lines_end": 1, + } + return yaml.safe_dump({"code_suggestions": [suggestion]}), "stop" + + async def reflect(*args): + return "" + + tool.ai_handler = SimpleNamespace(chat_completion=completion) + tool._self_reflect_with_fallback = reflect + monkeypatch.setattr(module, "get_pr_multi_diffs", lambda *a, **k: ["a", "b", "c"]) + return tool, calls + + +@pytest.mark.parametrize("parallel", [True, False]) +async def test_recovery_restores_missing_results_without_replacing_successes(configured, monkeypatch, parallel): + get_settings().set("pr_code_suggestions.parallel_calls", parallel) + tool, calls = make_tool(monkeypatch, {("gpt-4o", "b"): TimeoutError("provider timeout")}) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py", "c.py"] + assert [(m, chunk, dep) for m, chunk, dep, _ in calls] == [ + ("gpt-4o", "a", "primary"), ("gpt-4o", "b", "primary"), + ("gpt-4o", "c", "primary"), ("gpt-4o-mini", "b", "secondary"), + ] + assert tool.failed_chunk_count == 0 + assert tool.total_chunk_count == 3 + assert get_settings().get("openai.deployment_id") == "primary" + + +async def test_recovery_exhaustion_keeps_successes_and_reports_remaining_gap(configured, monkeypatch): + tool, calls = make_tool(monkeypatch, { + ("gpt-4o", "b"): RuntimeError("primary failed"), + ("gpt-4o", "c"): RuntimeError("primary failed"), + ("gpt-4o-mini", "b"): RuntimeError("secondary failed"), + ("gpt-4o-mini", "c"): RuntimeError("secondary failed"), + ("gpt-4.1", "c"): RuntimeError("last failed"), + }) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py"] + assert tool.failed_chunk_count == 1 + assert [(m, c) for m, c, _, _ in calls].count(("gpt-4o", "a")) == 1 + assert [c for m, c, _, _ in calls if m == "gpt-4.1"] == ["b", "c"] + + +async def test_healthy_run_never_enters_recovery(configured, monkeypatch): + tool, calls = make_tool(monkeypatch, {}) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert len(result["code_suggestions"]) == 3 + assert tool.failed_chunk_count == 0 + assert {m for m, _, _, _ in calls} == {"gpt-4o"} + + +@pytest.mark.parametrize("fallbacks, expected", [([], ["a.py", "c.py"]), (["gpt-4o-mini"], ["a.py", "b.py", "c.py"])]) +async def test_fallback_recovery_runs_whenever_fallbacks_are_set(configured, monkeypatch, fallbacks, expected): + # Recovery is the default behaviour whenever config.fallback_models is set, + # so an empty fallback list preserves the partial-success policy while a + # configured fallback repairs the failed chunk. + get_settings().set("config.fallback_models", fallbacks) + tool, calls = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == expected + assert tool.failed_chunk_count == (1 if not fallbacks else 0) + assert len(calls) == (3 if not fallbacks else 4) + + +async def test_larger_fallback_recovers_when_earlier_one_is_over_budget(configured, monkeypatch): + # The first fallback model cannot fit the full prompt, so it is skipped and + # recovery still moves on to the later model instead of giving up. + get_settings().set("config.fallback_models", ["gpt-4o-mini", "gpt-4.1"]) + monkeypatch.setattr(module, "get_max_tokens", lambda model: 1600 if model == "gpt-4o-mini" else 10000) + tool, calls = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + tool.vars["instructions"] = "must retain these instructions " * 500 + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py", "c.py"] + assert tool.failed_chunk_count == 0 + assert [(m, c) for m, c, _, _ in calls if m == "gpt-4.1"] == [("gpt-4.1", "b")] + + +async def test_all_failed_primary_keeps_existing_outer_fallback(configured, monkeypatch): + tool, calls = make_tool(monkeypatch, {("gpt-4o", c): RuntimeError("failure") for c in "abc"}) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py", "c.py"] + assert [m for m, _, _, _ in calls] == ["gpt-4o"] * 3 + ["gpt-4o-mini"] * 3 + assert tool.failed_chunk_count == 0 + + +async def test_partial_success_on_outer_fallback_only_tries_later_models(configured, monkeypatch): + failures = {("gpt-4o", c): RuntimeError("failure") for c in "abc"} + failures[("gpt-4o-mini", "b")] = RuntimeError("secondary failed") + tool, calls = make_tool(monkeypatch, failures) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py", "c.py"] + assert [m for m, _, _, _ in calls] == ["gpt-4o"] * 3 + ["gpt-4o-mini"] * 3 + ["gpt-4.1"] + assert get_settings().get("openai.deployment_id") == "primary" + + +@pytest.mark.parametrize("remaining_model", [True, False]) +async def test_oversized_fallback_is_skipped_without_truncating_context(configured, monkeypatch, remaining_model): + if not remaining_model: + get_settings().set("config.fallback_models", ["gpt-4o-mini"]) + tool, calls = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + tool.vars["instructions"] = "must retain these instructions " * 500 + # Use the real per-model tokenizer, with a deliberately smaller configured capacity. + monkeypatch.setattr(module, "get_max_tokens", lambda model: 1600 if model == "gpt-4o-mini" else 10000) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert not any(m == "gpt-4o-mini" for m, _, _, _ in calls) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ( + ["a.py", "b.py", "c.py"] if remaining_model else ["a.py", "c.py"] + ) + assert tool.failed_chunk_count == (0 if remaining_model else 1) + assert all(system == "Review " + tool.vars["instructions"] for _, _, _, system in calls) + + +async def test_marker_text_in_diff_is_counted_literally_for_recovery(configured, monkeypatch): + tool, calls = make_tool(monkeypatch, {("gpt-4o", "<|endoftext|>"): RuntimeError("failure")}) + # Patch after make_tool so the fixture's default chunk list does not win. + monkeypatch.setattr(module, "get_pr_multi_diffs", lambda *a, **k: ["a", "<|endoftext|>", "c"]) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert ("gpt-4o-mini", "<|endoftext|>", "secondary") in [(m, c, d) for m, c, d, _ in calls] + assert len(result["code_suggestions"]) == 3 + + +async def test_boundary_size_chunk_is_eligible_for_recovery(configured, monkeypatch): + tool, calls = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + monkeypatch.setattr(module.TokenHandler, "count_tokens", lambda self, s: 1) + monkeypatch.setattr(module, "get_max_tokens", lambda model: 1502 if model == "gpt-4o-mini" else 10000) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py", "c.py"] + assert tool.failed_chunk_count == 0 + + +async def test_recovery_reserves_handler_output_budget_for_fallback(configured, monkeypatch): + # When the AI handler exposes a concrete output allowance (e.g. a large + # config.max_output_tokens), recovery reserves it for the completion instead of + # the fixed soft threshold. The first fallback then cannot fit the prompt + # (5001 - 5000 = 1 < 2 tokens), while the fixed 1500 threshold would have let + # it recover the chunk; only the later model recovers it. + tool, calls = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + reserves = [] + + def get_output_token_reserve(model, default): + reserves.append((model, get_settings().get("openai.deployment_id"))) + return 5000 + + tool.ai_handler = SimpleNamespace( + chat_completion=tool.ai_handler.chat_completion, + get_output_token_reserve=get_output_token_reserve, + ) + monkeypatch.setattr(module.TokenHandler, "count_tokens", lambda self, s: 1) + monkeypatch.setattr(module, "get_max_tokens", lambda model: 5001 if model == "gpt-4o-mini" else 10000) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py", "c.py"] + assert tool.failed_chunk_count == 0 + assert [m for m, _, _, _ in calls] == ["gpt-4o"] * 3 + ["gpt-4.1"] + assert [(m, c) for m, c, _, _ in calls if m == "gpt-4.1"] == [("gpt-4.1", "b")] + assert reserves == [("gpt-4o-mini", "secondary"), ("gpt-4.1", "last")] + + +async def test_empty_prediction_is_a_success_not_a_retry_trigger(configured, monkeypatch): + tool, calls = make_tool(monkeypatch, { + ("gpt-4o", "a"): "code_suggestions: []", + ("gpt-4o", "b"): TimeoutError("failure"), + ("gpt-4o-mini", "b"): "code_suggestions: []", + }) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["c.py"] + assert tool.failed_chunk_count == 0 + assert len(calls) == 4 + + +async def test_invalid_yaml_recovery_preserves_parse_failure_notice(configured, monkeypatch): + tool, calls = make_tool(monkeypatch, { + ("gpt-4o", "b"): TimeoutError("failure"), + ("gpt-4o-mini", "b"): "not a suggestions mapping", + }) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "c.py"] + assert tool.failed_chunk_count == 1 + assert tool.parse_failure_count == 1 + assert len(calls) == 4 + + +@pytest.mark.parametrize("parent_cancel", [True, False]) +async def test_cancellation_stops_recovery_and_restores_deployment(configured, monkeypatch, parent_cancel): + started = asyncio.Event() + finished = asyncio.Event() + + async def wait_or_cancel(): + started.set() + try: + if parent_cancel: + await asyncio.Event().wait() + raise asyncio.CancelledError() + finally: + finished.set() + + tool, calls = make_tool(monkeypatch, { + ("gpt-4o", "b"): RuntimeError("failure"), + ("gpt-4o-mini", "b"): wait_or_cancel, + }) + task = asyncio.create_task(retry_with_fallback_models(tool.prepare_prediction_main)) + await asyncio.wait_for(started.wait(), timeout=5) + if parent_cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert finished.is_set() + assert not any(m == "gpt-4.1" for m, _, _, _ in calls) + assert get_settings().get("openai.deployment_id") == "primary" + tool.git_provider.publish_comment.assert_not_called() + + +async def test_model_switch_waits_for_successful_sibling_to_finish(configured, monkeypatch): + completed = asyncio.Event() + + async def finish_later(): + await asyncio.sleep(0) + assert get_settings().get("openai.deployment_id") == "primary" + completed.set() + + async def recover(): + assert completed.is_set() + assert get_settings().get("openai.deployment_id") == "secondary" + + tool, _ = make_tool(monkeypatch, { + ("gpt-4o", "a"): finish_later, + ("gpt-4o", "b"): RuntimeError("failure"), + ("gpt-4o-mini", "b"): recover, + }) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert len(result["code_suggestions"]) == 3 + + +async def test_instances_do_not_reuse_predictions_or_prompt_context(configured, monkeypatch): + first, first_calls = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + await retry_with_fallback_models(first.prepare_prediction_main) + second, second_calls = make_tool(monkeypatch, {("gpt-4o", "a"): RuntimeError("failure")}) + second.vars["instructions"] = "scope-two" + result = await retry_with_fallback_models(second.prepare_prediction_main) + assert len(result["code_suggestions"]) == 3 + assert all(system == "Review scope-one" for _, _, _, system in first_calls) + assert all(system == "Review scope-two" for _, _, _, system in second_calls) + assert [c for m, c, _, _ in second_calls if m == "gpt-4o-mini"] == ["a"] + + +async def test_concurrent_requests_keep_separate_deployments(configured, monkeypatch): + both_recovering = asyncio.Event() + started = [] + + async def invoke(name): + with request_cycle_context({"settings": copy.deepcopy(get_settings())}): + get_settings().set("openai.deployment_id", name + "-primary") + get_settings().set("openai.fallback_deployments", [name + "-secondary", name + "-last"]) + + async def overlap(): + started.append(name) + if len(started) == 2: + both_recovering.set() + await asyncio.wait_for(both_recovering.wait(), timeout=5) + assert get_settings().get("openai.deployment_id") == name + "-secondary" + + tool, calls = make_tool(monkeypatch, { + ("gpt-4o", "b"): TimeoutError("failure"), ("gpt-4o-mini", "b"): overlap, + }) + tool.vars["instructions"] = name + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert len(result["code_suggestions"]) == 3 + assert all(dep.startswith(name) and system == "Review " + name for _, _, dep, system in calls) + assert get_settings().get("openai.deployment_id") == name + "-primary" + + await asyncio.gather(invoke("one"), invoke("two")) + assert get_settings().get("openai.deployment_id") == "primary" + + +@pytest.mark.parametrize("recovered", [True, False]) +async def test_published_suggestions_and_coverage_match_completed_chunks(configured, monkeypatch, recovered): + failures = {("gpt-4o", "b"): RuntimeError("primary failed")} + if not recovered: + failures.update({(model, "b"): RuntimeError("failed") for model in ["gpt-4o-mini", "gpt-4.1"]}) + tool, _ = make_tool(monkeypatch, failures) + tool.git_provider.diff_files = [FilePatchInfo("old()\n", "old()\n", "", c + ".py") for c in "abc"] + tool.git_provider.publish_code_suggestions.return_value = True + tool.git_provider.supports_code_suggestions_artifact.return_value = False + init_run_details() + result = await retry_with_fallback_models(tool.prepare_prediction_main) + await tool.push_inline_code_suggestions(result) + published = tool.git_provider.publish_code_suggestions.call_args.args[0] + assert [s["relevant_file"] for s in published] == (["a.py", "b.py", "c.py"] if recovered else ["a.py", "c.py"]) + assert all("```suggestion\nnew()\n```" in s["body"] for s in published) + assert bool(tool._get_suggestions_coverage_footer()) is not recovered + assert tool.git_provider.publish_comment.called is not recovered + details = get_run_details() + assert details.fallback_used is recovered + # The outer wrapper records the primary model after recovery restores the deployment, + # so the model line names the primary while the sticky flag preserves the fallback. + assert details.model_used == "gpt-4o" + assert details.num_ai_calls == 0 + + +async def test_recovered_suggestion_still_passes_anchor_validation(configured, monkeypatch): + tool, _ = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + tool.git_provider.diff_files = [FilePatchInfo("", "old()\n", "", c + ".py") for c in "ac"] + tool.git_provider.supports_code_suggestions_artifact.return_value = False + result = await retry_with_fallback_models(tool.prepare_prediction_main) + await tool.push_inline_code_suggestions(result) + published = tool.git_provider.publish_code_suggestions.call_args.args[0] + assert [s["relevant_file"] for s in published] == ["a.py", "c.py"] + assert "b.py" in tool.git_provider.publish_comment.call_args.args[0] + assert "```suggestion" not in tool.git_provider.publish_comment.call_args.args[0] + + +async def test_recovered_run_details_report_primary_model_and_sticky_fallback(configured, monkeypatch): + tool, _ = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + tool.git_provider.diff_files = [FilePatchInfo("old()\n", "old()\n", "", c + ".py") for c in "abc"] + tool.git_provider.publish_code_suggestions.return_value = True + tool.git_provider.supports_code_suggestions_artifact.return_value = False + init_run_details() + result = await retry_with_fallback_models(tool.prepare_prediction_main) + await tool.push_inline_code_suggestions(result) + details = get_run_details() + assert details.fallback_used is True + assert details.model_used == "gpt-4o" + assert details.num_ai_calls == 0 + + +async def test_ambiguous_model_deployment_chain_keeps_partial_result(configured, monkeypatch): + get_settings().set("config.fallback_models", ["gpt-4o", "gpt-4.1"]) + get_settings().set("openai.fallback_deployments", ["primary", "last"]) + tool, calls = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "c.py"] + assert len(calls) == 3 + + +async def test_repeated_fallback_pair_skips_recoveries_without_retrying_it(configured, monkeypatch): + # A unique fallback pair appearing twice in the chain would route the same + # failed chunk to an identical (model, deployment) a second time, so recovery + # gives up and keeps the partial result instead of adding duplicate inference. + get_settings().set("config.fallback_models", ["gpt-4o-mini", "gpt-4o-mini"]) + get_settings().set("openai.fallback_deployments", ["secondary", "secondary"]) + tool, calls = make_tool(monkeypatch, {("gpt-4o", "b"): RuntimeError("failure")}) + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "c.py"] + assert len(calls) == 3 + assert tool.failed_chunk_count == 1 + + +async def test_recovery_uses_existing_reflection_before_publishing_results(configured, monkeypatch): + keys = ("config.model_reasoning", "pr_code_suggestions_reflect_prompt.system", + "pr_code_suggestions_reflect_prompt.user") + snapshot = snapshot_settings(keys) + try: + get_settings().set(keys[0], "gpt-4o") + get_settings().set(keys[1], "Reflect") + get_settings().set(keys[2], "{{ diff }}") + tool, _ = make_tool(monkeypatch, {("gpt-4o", "b"): TimeoutError("failure")}) + del tool._self_reflect_with_fallback + generation = tool.ai_handler.chat_completion + reflections = [] + + async def completion(*, model, system, user, **kwargs): + if system != "Reflect": + return await generation(model=model, system=system, user=user, **kwargs) + reflections.append((model, user, get_settings().get("openai.deployment_id"))) + return yaml.safe_dump({"code_suggestions": [{"suggestion_score": 9, "why": "Verified"}]}), "stop" + + tool.ai_handler.chat_completion = completion + result = await retry_with_fallback_models(tool.prepare_prediction_main) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py", "c.py"] + assert all(s["score"] == 9 and s["score_why"] == "Verified" for s in result["code_suggestions"]) + assert ("gpt-4o-mini", "b", "secondary") in reflections + assert len(reflections) == 3 + finally: + restore_settings(snapshot) + + +async def test_routed_primary_recovers_failed_slots_with_configured_fallback(configured, monkeypatch): + # A routed primary ([model_routing]) replaces config.model ahead of the same + # fallbacks; recovery must reproduce that substitution instead of scanning + # the configured chain. The routed model is kept out of the fallback list so + # the (model, deployment) pair stays unique in the effective chain. + get_settings().set("config.fallback_models", ["gpt-4.1"]) + get_settings().set("openai.fallback_deployments", ["last"]) + get_settings().set("model_routing.enable", True) + get_settings().set("model_routing.rules", + [{"model": "gpt-4o-mini", "max_files": 10, "deployment_id": "secondary"}]) + tool, calls = make_tool(monkeypatch, {("gpt-4o-mini", "b"): RuntimeError("failure")}) + tool.git_provider.get_diff_files.return_value = [FilePatchInfo("old()\n", "old()\n", "", c + ".py") for c in "abc"] + result = await retry_with_fallback_models(tool.prepare_prediction_main, + git_provider=tool.git_provider) + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py", "c.py"] + assert [m for m, _, _, _ in calls] == ["gpt-4o-mini"] * 3 + ["gpt-4.1"] + assert get_settings().get("openai.deployment_id") == "primary" + tool.git_provider.get_diff_files.assert_called_once_with() + + +async def test_routing_is_not_repeated_during_recovery(configured, monkeypatch): + get_settings().set("config.fallback_models", ["gpt-4.1"]) + get_settings().set("openai.fallback_deployments", ["last"]) + get_settings().set("model_routing.enable", True) + get_settings().set("model_routing.rules", + [{"model": "gpt-4o-mini", "max_files": 10, "deployment_id": "secondary"}]) + tool, calls = make_tool(monkeypatch, {("gpt-4o-mini", "b"): RuntimeError("failure")}) + tool.git_provider.get_diff_files.side_effect = [ + [FilePatchInfo("old()\n", "old()\n", "", c + ".py") for c in "abc"], + RuntimeError("transient routing failure"), + ] + + result = await retry_with_fallback_models(tool.prepare_prediction_main, + git_provider=tool.git_provider) + + assert [s["relevant_file"] for s in result["code_suggestions"]] == ["a.py", "b.py", "c.py"] + assert [m for m, _, _, _ in calls] == ["gpt-4o-mini"] * 3 + ["gpt-4.1"] + tool.git_provider.get_diff_files.assert_called_once_with() diff --git a/tests/unittest/test_telemetry_config.py b/tests/unittest/test_telemetry_config.py index 8f1ada52e2..5c9eafd8d2 100644 --- a/tests/unittest/test_telemetry_config.py +++ b/tests/unittest/test_telemetry_config.py @@ -8,6 +8,7 @@ configuration when run through the real validation logic. """ +import os import tomllib from pathlib import Path @@ -143,6 +144,36 @@ def test_otlp_timeout_defaults_and_reads_setting(monkeypatch): assert get_otel_config().otlp_timeout == 10 +def test_prometheus_settings_provision_state_dir(monkeypatch, tmp_path): + state_dir = tmp_path / "prometheus-state" + _use_settings(monkeypatch, { + **VALID_ENABLED_SETTINGS, + "OTEL.EXPORTER_TYPE": "prometheus", + "OTEL.PROMETHEUS_MULTIPROC_DIR": str(state_dir), + }) + + config = get_otel_config() + + assert config.is_enabled is True + assert config.exporter_type == ExporterType.PROMETHEUS + assert config.prometheus_multiproc_dir == str(state_dir) + assert os.environ["PROMETHEUS_MULTIPROC_DIR"] == str(state_dir) + assert state_dir.is_dir(), "the multiprocess state dir must be created eagerly" + + +def test_prometheus_without_endpoint_is_not_fail_closed(monkeypatch): + """Prometheus is a local scrape endpoint the operator explicitly selected, + so it must not require an OTLP endpoint like the 'otlp' exporter does.""" + _use_settings(monkeypatch, {**VALID_ENABLED_SETTINGS, "OTEL.EXPORTER_TYPE": "prometheus"}) + + with capture_loguru(level="WARNING") as captured: + config = get_otel_config() + + assert config.is_enabled is True + assert config.exporter_type == ExporterType.PROMETHEUS + assert captured == [], "prometheus must not emit a fail-closed warning" + + def test_otlp_protocol_defaults_to_http_and_reads_setting(monkeypatch): _use_settings(monkeypatch, dict(VALID_ENABLED_SETTINGS)) assert get_otel_config().otlp_protocol == "http" @@ -230,6 +261,7 @@ def test_shipped_configuration_toml_otel_section_values(): assert configuration["environment"] == "development" assert configuration["otlp_timeout"] == 3 assert configuration["otlp_protocol"] == "http", "the default transport must not need an extra" + assert configuration["prometheus_multiproc_dir"] == "/tmp/pr-agent-prometheus" assert configuration["include_pr_url"] is False, "PR URLs must be opt-in (privacy)" assert configuration["include_error_details"] is False, "error details must be opt-in (privacy)" diff --git a/tests/unittest/test_telemetry_exporter_types.py b/tests/unittest/test_telemetry_exporter_types.py index 85fdb00ecc..3fb766e599 100644 --- a/tests/unittest/test_telemetry_exporter_types.py +++ b/tests/unittest/test_telemetry_exporter_types.py @@ -32,7 +32,9 @@ def _reset_telemetry(): def test_valid_exporter_types_equals_constant_set(): - assert VALID_EXPORTER_TYPES == {ExporterType.CONSOLE, ExporterType.OTLP, ExporterType.NONE} + assert VALID_EXPORTER_TYPES == { + ExporterType.CONSOLE, ExporterType.OTLP, ExporterType.PROMETHEUS, ExporterType.NONE + } def test_disabled_tracer_is_noop_never_global(monkeypatch): @@ -64,6 +66,7 @@ def test_exporter_type_constant_values(): users put in configuration.toml.""" assert ExporterType.OTLP == "otlp" assert ExporterType.CONSOLE == "console" + assert ExporterType.PROMETHEUS == "prometheus" assert ExporterType.NONE == "none" @@ -85,6 +88,22 @@ def test_metric_exporter_none(): assert _create_metric_exporter(make_config(exporter_type=ExporterType.NONE)) is None +def test_span_exporter_prometheus_is_none(): + """The prometheus exporter is metrics-only: spans are silently dropped.""" + with capture_loguru(level="WARNING") as captured: + exporter = _create_exporter(make_config(exporter_type=ExporterType.PROMETHEUS)) + + assert exporter is None + assert "metrics only" in "\n".join(captured) + + +def test_metric_exporter_prometheus(): + from pr_agent.telemetry.prometheus import PrometheusMetricExporter + + exporter = _create_metric_exporter(make_config(exporter_type=ExporterType.PROMETHEUS)) + assert isinstance(exporter, PrometheusMetricExporter) + + def test_span_exporter_unknown_returns_none(): assert _create_exporter(make_config(exporter_type="bogus")) is None diff --git a/tests/unittest/test_telemetry_prometheus.py b/tests/unittest/test_telemetry_prometheus.py new file mode 100644 index 0000000000..3d6ecab7f1 --- /dev/null +++ b/tests/unittest/test_telemetry_prometheus.py @@ -0,0 +1,129 @@ +"""Tests for the native Prometheus exporter and its `/metrics` endpoint. + +The whole module runs with ``PROMETHEUS_MULTIPROC_DIR`` set (multiprocess +mode), mirroring a gunicorn worker: prometheus_client decides its storage mode +at import time from that env var, and the exporter defers the import to first +use. Nothing else in the telemetry suite imports prometheus_client, so the +env var above drives every process-global decision here. +""" + +import os +import subprocess +import sys + +import pytest +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + +from pr_agent.telemetry.prometheus import ( + PrometheusMetricExporter, + attach_metrics_endpoint, + prometheus_metrics_enabled, + prometheus_response, +) +from pr_agent.telemetry.prometheus_multiproc import ensure_prometheus_multiproc_dir + + +@pytest.fixture(scope="session", autouse=True) +def _multiproc_env(tmp_path_factory): + """Set the multiprocess state dir before the first prometheus_client import.""" + state_dir = tmp_path_factory.mktemp("prometheus-multiproc") + ensure_prometheus_multiproc_dir(str(state_dir)) + yield str(state_dir) + + +def _counter_provider(): + reader = PeriodicExportingMetricReader(PrometheusMetricExporter()) + provider = MeterProvider(metric_readers=[reader]) + return provider, reader + + +def test_counter_export_renders_prometheus_text(): + provider, reader = _counter_provider() + counter = provider.get_meter("prometheus-test").create_counter( + "pr_agent.commands", unit="{command}", description="PR-Agent commands executed" + ) + counter.add(3, {"pr_agent.command": "review", "provider.name": "github"}) + reader.collect() + + body = prometheus_response().body.decode() + # The OTel name "pr_agent.commands" is sanitized (dots -> underscores) and + # Prometheus appends the counter's _total suffix. + assert "# TYPE pr_agent_commands_total counter" in body + assert "# HELP pr_agent_commands_total PR-Agent commands executed" in body + assert 'pr_agent_commands_total{pr_agent_command="review",provider_name="github"} 3.0' in body + + +def test_multiple_workers_aggregate_at_scrape_time(): + """Run the exporter in two fresh interpreter processes (as gunicorn workers + would), each recording one command, and assert /metrics merges them.""" + child_code = "; ".join([ + "from opentelemetry.sdk.metrics import MeterProvider", + "from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader", + "from pr_agent.telemetry.prometheus import PrometheusMetricExporter", + "provider = MeterProvider(metric_readers=[PeriodicExportingMetricReader(PrometheusMetricExporter())])", + "counter = provider.get_meter('mp-test').create_counter('pr_agent.commands', unit='{command}', description='PR-Agent commands executed')", + "counter.add(1, {'pr_agent.command': 'review'})", + "provider.shutdown()", + ]) + for _ in range(2): + subprocess.run([sys.executable, "-c", child_code], check=True, capture_output=True) + + body = prometheus_response().body.decode() + # Both workers' files are merged into one series at scrape time; the + # parent's own in-process contribution (different label set) also shows. + assert 'pr_agent_commands_total{pr_agent_command="review"} 2.0' in body + assert 'pr_agent_command="review",provider_name="github"}' in body + + +def test_attach_metrics_endpoint_registers_lazy_get_route(): + from fastapi import APIRouter + + router = APIRouter() + attach_metrics_endpoint(router) + + assert len(router.routes) == 1 + route = router.routes[0] + assert route.path == "/metrics" + assert "GET" in route.methods + assert route.include_in_schema is False + # The endpoint resolves prometheus_client only on request — registering the + # route must not pull it in (a gunicorn master importing the app would + # otherwise lock in the wrong multiprocess mode before `when_ready`). + assert callable(route.endpoint) + + +def test_prometheus_metrics_enabled_reads_settings(monkeypatch): + from pr_agent.telemetry import prometheus as prometheus_module + + class FakeSettings: + def __init__(self, values): + self._values = values + + def get(self, key, default=None): + return self._values.get(key, default) + + monkeypatch.setattr(prometheus_module, "get_settings", lambda: FakeSettings({})) + assert prometheus_metrics_enabled() is False + + monkeypatch.setattr(prometheus_module, "get_settings", lambda: FakeSettings( + {"OTEL.IS_ENABLED": True, "OTEL.EXPORTER_TYPE": "prometheus"})) + assert prometheus_metrics_enabled() is True + + monkeypatch.setattr(prometheus_module, "get_settings", lambda: FakeSettings( + {"OTEL.IS_ENABLED": True, "OTEL.EXPORTER_TYPE": "otlp"})) + assert prometheus_metrics_enabled() is False + + +def test_importing_prometheus_bridge_does_not_import_prometheus_client(): + """The module must never import prometheus_client at import time — a gunicorn + master that imports the app under preload_app would otherwise decide the + non-multiprocess storage mode before when_ready provisions the state dir.""" + env = {k: v for k, v in os.environ.items() if k != "PROMETHEUS_MULTIPROC_DIR"} + code = ( + "import sys; " + "from pr_agent.telemetry.prometheus import prometheus_metrics_enabled, " + "attach_metrics_endpoint; " + "assert 'prometheus_client' not in sys.modules" + ) + subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, env=env) diff --git a/tests/unittest/test_telemetry_prometheus_multiproc.py b/tests/unittest/test_telemetry_prometheus_multiproc.py new file mode 100644 index 0000000000..146199876c --- /dev/null +++ b/tests/unittest/test_telemetry_prometheus_multiproc.py @@ -0,0 +1,59 @@ +"""Tests for the prometheus multiprocess state directory provisioning. + +The state directory is the one piece of the exporter the gunicorn master +touches before workers import prometheus_client, and a predictable default in +a world-writable prefix like /tmp must not be redirectable to a foreign +location or shared with other users. These tests lock in the symlink and +ownership refusals and the 0700 creation. +""" + +import os +import stat + +import pytest + +from pr_agent.telemetry.prometheus_multiproc import ( + PROMETHEUS_MULTIPROC_DIR_ENV, + ensure_prometheus_multiproc_dir, +) + + +def test_fresh_dir_created_0700_and_env_set(tmp_path, monkeypatch): + monkeypatch.setenv(PROMETHEUS_MULTIPROC_DIR_ENV, "stale") + state_dir = tmp_path / "prometheus-state" + + result = ensure_prometheus_multiproc_dir(str(state_dir)) + + assert result == str(state_dir) + assert os.environ[PROMETHEUS_MULTIPROC_DIR_ENV] == str(state_dir) + assert state_dir.is_dir() + assert stat.S_IMODE(state_dir.stat().st_mode) == 0o700 + + +def test_existing_owned_dir_accepted_and_tightened_to_0700(tmp_path): + state_dir = tmp_path / "prometheus-state" + state_dir.mkdir(mode=0o755) + + ensure_prometheus_multiproc_dir(str(state_dir)) + + assert stat.S_IMODE(state_dir.stat().st_mode) == 0o700 + + +def test_symlink_is_refused(tmp_path): + real = tmp_path / "real-state" + real.mkdir() + link = tmp_path / "prometheus-state" + link.symlink_to(real) + + with pytest.raises(RuntimeError, match="symlinked"): + ensure_prometheus_multiproc_dir(str(link)) + + +def test_dir_owned_by_foreign_uid_is_refused(tmp_path, monkeypatch): + state_dir = tmp_path / "prometheus-state" + state_dir.mkdir() + foreign_uid = os.getuid() + 1 + monkeypatch.setattr("pr_agent.telemetry.prometheus_multiproc.os.getuid", lambda: foreign_uid) + + with pytest.raises(RuntimeError, match="owned by uid"): + ensure_prometheus_multiproc_dir(str(state_dir)) diff --git a/tests/unittest/test_ticket_compliance.py b/tests/unittest/test_ticket_compliance.py index 2fe8bbf1d8..30801e9462 100644 --- a/tests/unittest/test_ticket_compliance.py +++ b/tests/unittest/test_ticket_compliance.py @@ -6,6 +6,8 @@ - Authenticated task fetching and provider integration - Edge cases (mixed content, no tickets, duplicates, API failures) """ +import asyncio + import pytest from pr_agent.config_loader import get_settings @@ -360,6 +362,41 @@ async def test_extract_tickets_caps_non_github_asana_references(self): for ticket in tickets ) + @pytest.mark.asyncio + async def test_fallback_asana_tickets_count_against_jira_budget(self, monkeypatch): + """The fallback path shares MAX_TICKETS between Asana and Jira results, + matching the provider-specific paths that pass their existing tickets in.""" + jira_keys = ["jira.jira_site", "jira.jira_api_email", "jira.jira_api_token"] + snapshot = _settings_helpers.snapshot_settings(jira_keys) + try: + get_settings().set("jira.jira_site", "example") + get_settings().set("jira.jira_api_email", "user@example.com") + get_settings().set("jira.jira_api_token", "test-token") + + class _FakeJira: + def issue(self, key): + return {"fields": {"summary": f"Ticket {key}", "description": "body", + "labels": []}} + + monkeypatch.setattr(tpc, "Jira", lambda **kwargs: _FakeJira()) + provider = _GenericProvider( + "Related Asana tasks: " + "https://app.asana.com/0/99/111111111111 " + "https://app.asana.com/0/99/222222222222 " + "https://app.asana.com/0/99/333333333333 " + "Also PROJ-1 PROJ-2 PROJ-3 PROJ-4 PROJ-5" + ) + + tickets = await tpc.extract_tickets(provider) + + assert len(tickets) == tpc.MAX_TICKETS + assert [t["ticket_id"] for t in tickets[:3]] == [ + "111111111111", "222222222222", "333333333333" + ] + assert [t["ticket_id"] for t in tickets[3:]] == ["PROJ-1", "PROJ-2"] + finally: + _settings_helpers.restore_settings(snapshot) + @pytest.mark.asyncio async def test_extract_tickets_adds_asana_without_truncating_azure(self): """Azure work items remain intact while Asana additions keep their own cap.""" @@ -504,6 +541,152 @@ async def _failing_fetch(_session, ticket_url, _max_body_characters): assert tickets == [] assert attempted_urls == ticket_urls[:3] + @pytest.mark.asyncio + async def test_asana_ticket_fetches_run_concurrently_and_preserve_input_order(self, monkeypatch): + active_fetches = 0 + max_active_fetches = 0 + all_started = asyncio.Event() + + async def _coordinated_fetch(_session, ticket_url, _max_body_characters): + nonlocal active_fetches, max_active_fetches + active_fetches += 1 + max_active_fetches = max(max_active_fetches, active_fetches) + if active_fetches == 3: + all_started.set() + try: + await asyncio.wait_for(all_started.wait(), timeout=0.1) + task_gid = tpc._get_asana_task_gid(ticket_url) + return { + "ticket_id": task_gid, + "ticket_url": ticket_url, + "title": f"Asana task {task_gid}", + "body": "", + "labels": "", + } + finally: + active_fetches -= 1 + + monkeypatch.setattr(tpc, "_fetch_asana_ticket_content", _coordinated_fetch) + ticket_urls = [ + f"https://app.asana.com/0/99/{task_gid}" + for task_gid in (111111111111, 222222222222, 333333333333) + ] + + tickets = await tpc._fetch_asana_ticket_contents(ticket_urls, 3, 10000) + + assert max_active_fetches == 3 + assert [ticket["ticket_url"] for ticket in tickets] == ticket_urls + + @pytest.mark.asyncio + async def test_asana_ticket_fetches_keep_successes_ordered_after_partial_failure(self, monkeypatch): + async def _partially_failing_fetch(_session, ticket_url, _max_body_characters): + task_gid = tpc._get_asana_task_gid(ticket_url) + await asyncio.sleep({"111111111111": 0.02, "222222222222": 0, "333333333333": 0.01}[task_gid]) + if task_gid == "222222222222": + raise RuntimeError("task unavailable") + return { + "ticket_id": task_gid, + "ticket_url": ticket_url, + "title": f"Asana task {task_gid}", + "body": "", + "labels": "", + } + + monkeypatch.setattr(tpc, "_fetch_asana_ticket_content", _partially_failing_fetch) + ticket_urls = [ + f"https://app.asana.com/0/99/{task_gid}" + for task_gid in (111111111111, 222222222222, 333333333333) + ] + + tickets = await tpc._fetch_asana_ticket_contents(ticket_urls, 3, 10000) + + assert [ticket["ticket_id"] for ticket in tickets] == ["111111111111", "333333333333"] + + @pytest.mark.asyncio + async def test_asana_ticket_timeout_does_not_cancel_other_fetches(self, monkeypatch): + async def _timeout_one_fetch(_session, ticket_url, _max_body_characters): + task_gid = tpc._get_asana_task_gid(ticket_url) + if task_gid == "222222222222": + raise TimeoutError("task timed out") + return { + "ticket_id": task_gid, + "ticket_url": ticket_url, + "title": f"Asana task {task_gid}", + "body": "", + "labels": "", + } + + monkeypatch.setattr(tpc, "_fetch_asana_ticket_content", _timeout_one_fetch) + ticket_urls = [ + f"https://app.asana.com/0/99/{task_gid}" + for task_gid in (111111111111, 222222222222, 333333333333) + ] + + tickets = await tpc._fetch_asana_ticket_contents(ticket_urls, 3, 10000) + + assert [ticket["ticket_id"] for ticket in tickets] == ["111111111111", "333333333333"] + + @pytest.mark.asyncio + async def test_asana_ticket_batch_cancellation_cleans_up_in_flight_fetches(self, monkeypatch): + active_fetches = 0 + all_started = asyncio.Event() + + async def _blocking_fetch(_session, _ticket_url, _max_body_characters): + nonlocal active_fetches + active_fetches += 1 + if active_fetches == 3: + all_started.set() + try: + await asyncio.Event().wait() + finally: + active_fetches -= 1 + + monkeypatch.setattr(tpc, "_fetch_asana_ticket_content", _blocking_fetch) + ticket_urls = [ + f"https://app.asana.com/0/99/{task_gid}" + for task_gid in (111111111111, 222222222222, 333333333333) + ] + batch = asyncio.create_task(tpc._fetch_asana_ticket_contents(ticket_urls, 3, 10000)) + await asyncio.wait_for(all_started.wait(), timeout=0.1) + + batch.cancel() + with pytest.raises(asyncio.CancelledError): + await batch + + assert active_fetches == 0 + + @pytest.mark.asyncio + async def test_asana_ticket_child_abort_cancels_and_drains_siblings(self, monkeypatch): + class FetchAborted(BaseException): + pass + + active_fetches = 0 + all_started = asyncio.Event() + + async def _aborting_fetch(_session, ticket_url, _max_body_characters): + nonlocal active_fetches + active_fetches += 1 + if active_fetches == 3: + all_started.set() + try: + await asyncio.wait_for(all_started.wait(), timeout=0.1) + if tpc._get_asana_task_gid(ticket_url) == "111111111111": + raise FetchAborted("child fetch aborted") + await asyncio.Event().wait() + finally: + active_fetches -= 1 + + monkeypatch.setattr(tpc, "_fetch_asana_ticket_content", _aborting_fetch) + ticket_urls = [ + f"https://app.asana.com/0/99/{task_gid}" + for task_gid in (111111111111, 222222222222, 333333333333) + ] + + with pytest.raises(FetchAborted, match="child fetch aborted"): + await tpc._fetch_asana_ticket_contents(ticket_urls, 3, 10000) + + assert active_fetches == 0 + @pytest.mark.asyncio async def test_invalid_asana_url_does_not_abort_valid_batch_entry(self, monkeypatch): attempted_urls = [] diff --git a/tests/unittest/test_ticket_extraction_async.py b/tests/unittest/test_ticket_extraction_async.py index b222679546..6644e284ed 100644 --- a/tests/unittest/test_ticket_extraction_async.py +++ b/tests/unittest/test_ticket_extraction_async.py @@ -679,7 +679,7 @@ def test_issue_reference_is_mapped_to_ticket_context(self, reference, settings_s "labels": "bug, backend", } ] - provider.gl.projects.get.assert_called_once_with("group/repo") + provider.gl.projects.get.assert_called_once_with("group/repo", lazy=True) project.issues.get.assert_called_once_with(7) diff --git a/tests/unittest/test_ticket_pr_compliance_capabilities.py b/tests/unittest/test_ticket_pr_compliance_capabilities.py new file mode 100644 index 0000000000..eac3bf1e2a --- /dev/null +++ b/tests/unittest/test_ticket_pr_compliance_capabilities.py @@ -0,0 +1,385 @@ +from unittest.mock import MagicMock + +import pytest + +from pr_agent.git_providers import ( + AzureDevopsProvider, + BitbucketProvider, + BitbucketServerProvider, + CodeCommitProvider, + GerritProvider, + GiteaProvider, + GithubProvider, + GitLabProvider, + GitProvider, + LocalGitProvider, + PlainDiffGitProvider, +) +from pr_agent.tools.ticket_pr_compliance_check import _provider_supports, extract_tickets + + +class _BaseStubProvider(GitProvider): + def is_supported(self, capability: str) -> bool: + return False + + def get_files(self) -> list: + return [] + + def get_diff_files(self) -> list: + return [] + + def publish_description(self, pr_title: str, pr_body: str): + pass + + def publish_code_suggestions(self, code_suggestions: list) -> bool: + return False + + def get_languages(self): + return {} + + def get_pr_branch(self): + return "" + + def get_user_id(self): + return "" + + def get_pr_description_full(self) -> str: + return "" + + def get_repo_settings(self): + return b"" + + def publish_comment(self, pr_comment: str, is_temporary: bool = False): + pass + + def publish_inline_comment( + self, body: str, relevant_file: str, relevant_line_in_file: str, original_suggestion=None + ): + pass + + def publish_inline_comments(self, comments: list[dict]): + pass + + def remove_initial_comment(self): + pass + + def remove_comment(self, comment): + pass + + def get_issue_comments(self): + return [] + + def publish_labels(self, labels): + pass + + def get_pr_labels(self, update=False): + return [] + + def remove_reaction(self, issue_comment_id: int, reaction_id: int) -> bool: + return False + + def get_commit_messages(self) -> str: + return "" + + +class _CustomGitHubSubstitute(_BaseStubProvider): + def __init__(self, description="Fixes #101", branch="feature/102-branch"): + self.repo = "owner/repo" + self.base_url_html = "https://github.com" + self._description = description + self._branch = branch + self.github_client = MagicMock() + self.repo_obj = MagicMock() + + issue_101 = MagicMock(number=101, title="Issue 101", body="Body 101", labels=["bug"]) + issue_102 = MagicMock(number=102, title="Issue 102", body="Body 102", labels=["core"]) + self.repo_obj.get_issue.side_effect = lambda num: {101: issue_101, 102: issue_102}[num] + + def supports_issue_url_tickets(self) -> bool: + return True + + def get_user_description(self) -> str: + return self._description + + def get_pr_branch(self) -> str: + return self._branch + + def _parse_issue_url(self, url: str): + num = int(url.rsplit("/", 1)[-1]) + return self.repo, num + + def fetch_sub_issues(self, ticket_url: str) -> list: + return [] + + +class _CustomGitLabSubstitute(_BaseStubProvider): + def __init__(self, description="Resolves group/project#55"): + self.id_project = "group/project" + self.gitlab_url = "https://gitlab.example.com" + self._description = description + self.gl = MagicMock() + + issue = MagicMock( + iid=55, + web_url="https://gitlab.example.com/group/project/-/issues/55", + title="GitLab Issue 55", + description="GitLab Description 55", + labels=["backend", "p1"], + ) + project = MagicMock() + project.issues.get.return_value = issue + self.gl.projects.get.return_value = project + + def supports_issue_reference_tickets(self) -> bool: + return True + + def get_user_description(self) -> str: + return self._description + + +class _CustomAzureSubstitute(_BaseStubProvider): + def __init__(self, work_items=None): + self._work_items = ( + work_items + if work_items is not None + else [ + { + "id": 999, + "url": "https://dev.azure.com/org/proj/_workitems/edit/999", + "title": "Work Item 999", + "body": "Work item body 999", + "acceptance_criteria": "AC 999", + "labels": ["azure-label"], + } + ] + ) + + def supports_linked_work_item_tickets(self) -> bool: + return True + + def get_linked_work_items(self) -> list: + return self._work_items + + +class _IndependentDuckTypedSubstitute: + def __init__(self, capability: str): + self._capability = capability + self.repo = "duck/repo" + self.base_url_html = "https://github.com" + self.id_project = "duck/repo" + self.gitlab_url = "https://gitlab.com" + self.repo_obj = MagicMock() + self.gl = MagicMock() + + def supports_issue_url_tickets(self) -> bool: + return self._capability == "github" + + def supports_issue_reference_tickets(self) -> bool: + return self._capability == "gitlab" + + def supports_linked_work_item_tickets(self) -> bool: + return self._capability == "azure" + + def get_user_description(self) -> str: + if self._capability == "github": + return "Fixes #42" + if self._capability == "gitlab": + return "Fixes duck/repo#42" + return "" + + def get_pr_branch(self) -> str: + return "main" + + def _parse_issue_url(self, url: str): + return self.repo, 42 + + def fetch_sub_issues(self, ticket_url: str) -> list: + return [] + + def get_linked_work_items(self) -> list: + return [ + { + "id": 42, + "url": "https://dev.azure.com/duck/repo/_workitems/edit/42", + "title": "Azure 42", + "body": "Body 42", + "acceptance_criteria": "", + "labels": [], + } + ] + + +def test_base_git_provider_capabilities_default_false(): + provider = _BaseStubProvider() + assert provider.supports_issue_url_tickets() is False + assert provider.supports_issue_reference_tickets() is False + assert provider.supports_linked_work_item_tickets() is False + + +@pytest.mark.parametrize( + ("provider_cls", "expected_github", "expected_gitlab", "expected_azure"), + [ + (GithubProvider, True, False, False), + (GitLabProvider, False, True, False), + (AzureDevopsProvider, False, False, True), + (BitbucketProvider, False, False, False), + (BitbucketServerProvider, False, False, False), + (CodeCommitProvider, False, False, False), + (LocalGitProvider, False, False, False), + (GerritProvider, False, False, False), + (GiteaProvider, False, False, False), + (PlainDiffGitProvider, False, False, False), + ], +) +def test_concrete_providers_declare_expected_ticket_capabilities( + provider_cls, expected_github, expected_gitlab, expected_azure +): + provider = provider_cls.__new__(provider_cls) + assert provider.supports_issue_url_tickets() is expected_github + assert provider.supports_issue_reference_tickets() is expected_gitlab + assert provider.supports_linked_work_item_tickets() is expected_azure + + +@pytest.mark.asyncio +async def test_extract_tickets_routes_to_github_capability(): + provider = _CustomGitHubSubstitute() + tickets = await extract_tickets(provider) + + assert tickets is not None + assert len(tickets) == 2 + assert tickets[0]["ticket_id"] == 101 + assert tickets[0]["title"] == "Issue 101" + assert tickets[1]["ticket_id"] == 102 + assert tickets[1]["title"] == "Issue 102" + + +@pytest.mark.asyncio +async def test_extract_tickets_routes_to_gitlab_capability(): + provider = _CustomGitLabSubstitute() + tickets = await extract_tickets(provider) + + assert tickets is not None + assert len(tickets) == 1 + assert tickets[0]["ticket_id"] == 55 + assert tickets[0]["title"] == "GitLab Issue 55" + assert tickets[0]["body"] == "GitLab Description 55" + assert tickets[0]["labels"] == "backend, p1" + + +@pytest.mark.asyncio +async def test_extract_tickets_routes_to_azure_capability(): + provider = _CustomAzureSubstitute() + tickets = await extract_tickets(provider) + + assert tickets is not None + assert len(tickets) == 1 + assert tickets[0]["ticket_id"] == 999 + assert tickets[0]["title"] == "Work Item 999" + assert tickets[0]["body"] == "Work item body 999" + assert tickets[0]["requirements"] == "AC 999" + assert tickets[0]["labels"] == "azure-label" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["github", "gitlab", "azure"]) +async def test_extract_tickets_supports_duck_typed_substitutes(route): + provider = _IndependentDuckTypedSubstitute(route) + if route == "github": + provider.repo_obj.get_issue.return_value = MagicMock(number=42, title="GH 42", body="GH Body", labels=[]) + elif route == "gitlab": + mock_proj = MagicMock() + mock_proj.issues.get.return_value = MagicMock( + iid=42, web_url="https://gitlab.com/duck/repo/-/issues/42", title="GL 42", description="GL Body", labels=[] + ) + provider.gl.projects.get.return_value = mock_proj + + tickets = await extract_tickets(provider) + assert tickets is not None + assert len(tickets) == 1 + assert tickets[0]["ticket_id"] == 42 + + +@pytest.mark.asyncio +async def test_extract_tickets_unsupported_provider_returns_none_without_asana(): + provider = _BaseStubProvider() + provider.get_user_description = lambda: "No tickets here" + result = await extract_tickets(provider) + assert result is None + + +@pytest.mark.asyncio +async def test_extract_tickets_unsupported_provider_preserves_asana_results(monkeypatch): + provider = _BaseStubProvider() + provider.get_user_description = lambda: "Related: https://app.asana.com/0/99/123456789012" + + async def fake_fetch(urls, max_tickets, max_characters): + assert urls == ["https://app.asana.com/0/99/123456789012"] + return [{"ticket_id": "123456789012", "ticket_url": urls[0]}] + + monkeypatch.setattr("pr_agent.tools.ticket_pr_compliance_check._fetch_asana_ticket_contents", fake_fetch) + tickets = await extract_tickets(provider) + assert tickets == [{"ticket_id": "123456789012", "ticket_url": "https://app.asana.com/0/99/123456789012"}] + + +@pytest.mark.asyncio +async def test_extract_tickets_github_substitute_handles_issue_fetch_error(): + provider = _CustomGitHubSubstitute(description="Fixes #999", branch="") + provider.repo_obj.get_issue.side_effect = RuntimeError("Issue not found") + tickets = await extract_tickets(provider) + assert tickets == [] + + +@pytest.mark.asyncio +async def test_extract_tickets_gitlab_substitute_handles_project_error(): + provider = _CustomGitLabSubstitute() + provider.gl.projects.get.side_effect = RuntimeError("GitLab API error") + tickets = await extract_tickets(provider) + assert tickets == [] + + +@pytest.mark.asyncio +async def test_extract_tickets_azure_substitute_handles_item_error(): + provider = _CustomAzureSubstitute(work_items=[{"invalid": "data"}]) + tickets = await extract_tickets(provider) + assert len(tickets) == 1 + assert tickets[0]["ticket_id"] is None + + +@pytest.mark.asyncio +async def test_extract_tickets_catches_top_level_exception(): + provider = _CustomGitHubSubstitute() + provider.supports_issue_url_tickets = MagicMock(side_effect=RuntimeError("unexpected crash")) + tickets = await extract_tickets(provider) + assert tickets == [] + + +def test_provider_supports_raises_for_unknown_capability(): + provider = _BaseStubProvider() + with pytest.raises(AttributeError, match="unknown provider capability: 'supports_nonexistent_capability'"): + _provider_supports(provider, "supports_nonexistent_capability") + + +def test_provider_supports_returns_false_for_missing_method_on_duck_typed_provider(): + duck_typed_provider = object() + assert _provider_supports(duck_typed_provider, "supports_issue_url_tickets") is False + + +def test_provider_supports_returns_false_for_non_callable_attribute(): + class AttributeProvider: + supports_issue_url_tickets = True + + assert _provider_supports(AttributeProvider(), "supports_issue_url_tickets") is False + + +def test_provider_supports_reads_callable_on_base_or_duck_typed_provider(): + class TruthyProvider: + def supports_issue_url_tickets(self): + return 1 + + class FalsyProvider: + def supports_issue_url_tickets(self): + return 0 + + assert _provider_supports(TruthyProvider(), "supports_issue_url_tickets") is True + assert _provider_supports(FalsyProvider(), "supports_issue_url_tickets") is False + assert _provider_supports(_BaseStubProvider(), "supports_issue_url_tickets") is False diff --git a/tests/unittest/test_tool_capability_guard_contract.py b/tests/unittest/test_tool_capability_guard_contract.py new file mode 100644 index 0000000000..85cbe244a2 --- /dev/null +++ b/tests/unittest/test_tool_capability_guard_contract.py @@ -0,0 +1,164 @@ +"""Contract: no tool may reach a git-provider method its provider has declared unsupported. + +`test_git_provider_method_contract.py` pins what providers DO for a given capability; +this file pins what CALLERS in `pr_agent/tools/` do with the answer. A tool that calls a +capability-gated `self.git_provider.(...)` without checking `is_supported()` first +raises (or silently corrupts output) on any provider that declines that capability: #2932 +was exactly this, for `get_labels` in `pr_reviewer.py`. + +Scope is intentionally the ENCLOSING FUNCTION, not the enclosing class: a guard several +frames away (a flag cached in `__init__`, a raise before a later call) is real, but it is +also invisible to the next person editing the callee, which is the failure mode this test +exists to catch. Every current call site either already checks locally or was given a +local check by this change (see `pr_update_changelog.py::_push_changelog_update` and +`pr_reviewer.py::_get_user_answers`) precisely so this test can be scoped this way without +false positives. +""" +import ast +from pathlib import Path + +TOOLS_DIR = Path(__file__).resolve().parent.parent.parent / "pr_agent" / "tools" +PROVIDERS_DIR = Path(__file__).resolve().parent.parent.parent / "pr_agent" / "git_providers" + +# Base GitProvider methods that at least one provider declines (see +# `_provider_declinable_capabilities` below), mapped to the `is_supported()` capability +# string that guards each. This mapping is domain knowledge (e.g. "get_labels covers +# publish_labels and get_pr_labels") and cannot be derived from the providers alone, so +# unlike the capability set below it is maintained by hand. +GUARDED_METHODS = { + "get_pr_labels": "get_labels", + "publish_labels": "get_labels", + "publish_inline_comments": "publish_inline_comments", + "create_inline_comment": "create_inline_comment", + "get_issue_comments": "get_issue_comments", + "create_or_update_pr_file": "push_code", +} + + +def _tool_files(): + return sorted(TOOLS_DIR.glob("*.py")) + + +def _is_supported_capability_strings(node): + """Every literal string passed as the first arg to an `is_supported(...)` call + anywhere under `node` (a module, class or function node).""" + caps = set() + for inner in ast.walk(node): + if (isinstance(inner, ast.Call) + and isinstance(inner.func, ast.Attribute) + and inner.func.attr == "is_supported" + and inner.args + and isinstance(inner.args[0], ast.Constant) + and isinstance(inner.args[0].value, str)): + caps.add(inner.args[0].value) + return caps + + +def _is_git_provider(node): + """`self.git_provider`, or a bare `git_provider` name in a helper that takes the + provider as a parameter (the pattern `pr_code_suggestions.py` already uses).""" + if isinstance(node, ast.Name): + return node.id == "git_provider" + return (isinstance(node, ast.Attribute) + and node.attr == "git_provider" + and isinstance(node.value, ast.Name) + and node.value.id == "self") + + +def _git_provider_call_sites(func_node): + """Every `self.git_provider.(...)` or `git_provider.(...)` call + directly under a function node, keyed by method name.""" + calls = {} + for inner in ast.walk(func_node): + if (isinstance(inner, ast.Call) + and isinstance(inner.func, ast.Attribute) + and _is_git_provider(inner.func.value)): + calls.setdefault(inner.func.attr, []).append(inner) + return calls + + +def _function_nodes(tree): + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + yield node + + +def _unguarded_call_sites(): + """(file, function, method, lineno) for every capability-gated + `self.git_provider.(...)` call with no matching `is_supported()` check + anywhere in its own enclosing function. + + This is membership, not control flow: it does not check that the guard runs + before the call, sits in the branch that actually reaches it, or that its result + is used. A guard placed after the call, in an unrelated branch, or with its + result ignored still satisfies this check. It catches a capability check being + absent from the function entirely, which is the regression this test guards + against; it will not catch one being present but misplaced. + """ + hits = [] + for path in _tool_files(): + tree = ast.parse(path.read_text(), filename=str(path)) + for func_node in _function_nodes(tree): + guarded_caps = _is_supported_capability_strings(func_node) + calls = _git_provider_call_sites(func_node) + for method, cap in GUARDED_METHODS.items(): + for call in calls.get(method, []): + if cap not in guarded_caps: + hits.append((path.name, func_node.name, method, call.lineno)) + return hits + + +def _provider_declinable_capabilities(): + """Every capability string at least one provider's `is_supported()` special-cases. + + Parsed from `pr_agent/git_providers/*.py` rather than hardcoded, so a new provider + declining a new capability grows this set (and the typo check below) automatically + instead of it silently going stale. + + Collects any string literal compared inside `is_supported()`, without checking + which branch it feeds or whether that branch returns False. A capability named in + a comparison that always returns True, or in an unrelated comparison, is treated + as declinable the same as one a provider genuinely rejects. + """ + caps = set() + for path in sorted(PROVIDERS_DIR.glob("*.py")): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == "is_supported": + for inner in ast.walk(node): + if isinstance(inner, ast.Compare): + for comparator in inner.comparators: + if isinstance(comparator, ast.Constant) and isinstance(comparator.value, str): + caps.add(comparator.value) + elif isinstance(comparator, (ast.List, ast.Tuple, ast.Set)): + for elt in comparator.elts: + if isinstance(elt, ast.Constant) and isinstance(elt.value, str): + caps.add(elt.value) + return caps + + +def _is_supported_strings_used_in_tools(): + caps = set() + for path in _tool_files(): + tree = ast.parse(path.read_text(), filename=str(path)) + caps |= _is_supported_capability_strings(tree) + return caps + + +def test_every_capability_gated_provider_call_is_locally_guarded(): + hits = _unguarded_call_sites() + assert hits == [], ( + "self.git_provider.() reachable without checking is_supported() in the " + "same function first: " + + ", ".join(f"{f}:{fn}() calls {m}() at line {ln}" for f, fn, m, ln in hits) + ) + + +def test_every_is_supported_string_in_tools_is_a_capability_some_provider_declines(): + used = _is_supported_strings_used_in_tools() + declinable = _provider_declinable_capabilities() + typos = used - declinable + assert not typos, ( + f"is_supported() called in pr_agent/tools/ with a string no provider ever " + f"declines, so the check always passes: {typos}" + ) diff --git a/tests/unittest/test_webhook_logic_core.py b/tests/unittest/test_webhook_logic_core.py index 1e67514a14..2f1cf34341 100644 --- a/tests/unittest/test_webhook_logic_core.py +++ b/tests/unittest/test_webhook_logic_core.py @@ -722,6 +722,8 @@ async def reject_duplicate(key, **kwargs): async def perform_commands(*args): performed.append(args) + monkeypatch.setattr(gitea_app, "apply_repo_settings", lambda _url: None) + monkeypatch.setattr(gitea_app, "should_process_pr_logic", lambda _body: True) monkeypatch.setattr(gitea_app, "push_trigger_slot", reject_duplicate) monkeypatch.setattr(gitea_app, "_perform_commands_gitea", perform_commands) api_url = "https://gitea.example.com/org/repo/pulls/1" @@ -736,6 +738,169 @@ async def perform_commands(*args): assert slots == [(api_url, {"allow_backlog": True, "ttl": 300})] +@pytest.mark.asyncio +async def test_gitea_push_applies_repo_settings_before_effective_gate(monkeypatch): + settings = get_settings() + original_gitea = copy.deepcopy(settings.get("GITEA")) + original_is_auto_command = settings.get("CONFIG.IS_AUTO_COMMAND") + settings.set("GITEA.HANDLE_PUSH_TRIGGER", False) + settings.set("GITEA.PUSH_COMMANDS", []) + calls = [] + + def apply_repo_settings(_url): + calls.append("settings") + get_settings().set("GITEA.HANDLE_PUSH_TRIGGER", True) + get_settings().set("GITEA.PUSH_COMMANDS", ["/review"]) + + def should_process_pr_logic(_body): + calls.append("filter") + return True + + @asynccontextmanager + async def record_slot(_key, **_kwargs): + calls.append("slot") + yield True + + class Agent: + async def handle_request(self, _url, command): + calls.append(command) + + monkeypatch.setattr(gitea_app, "apply_repo_settings", apply_repo_settings) + monkeypatch.setattr(gitea_app, "should_process_pr_logic", should_process_pr_logic) + monkeypatch.setattr(gitea_app, "push_trigger_slot", record_slot) + monkeypatch.setattr(gitea_app, "prepare_command", lambda command: command) + api_url = "https://gitea.example.com/org/repo/pulls/1" + body = { + "pull_request": {"url": api_url}, + "repository": {"full_name": "org/repo"}, + } + + try: + await gitea_app.handle_pr_event(body, "pull_request", "synchronized", Agent()) + finally: + settings.set("GITEA", original_gitea) + settings.set("CONFIG.IS_AUTO_COMMAND", original_is_auto_command) + + assert calls == ["settings", "filter", "slot", "/review"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("repo_trigger", "repo_commands"), + [ + (False, ["/review"]), + (True, []), + ], +) +async def test_gitea_effective_push_config_skips_before_reserving_slot( + monkeypatch, repo_trigger, repo_commands +): + settings = get_settings() + original_gitea = copy.deepcopy(settings.get("GITEA")) + settings.set("GITEA.HANDLE_PUSH_TRIGGER", True) + settings.set("GITEA.PUSH_COMMANDS", ["/host-review"]) + calls = [] + + def apply_repo_settings(_url): + calls.append("settings") + get_settings().set("GITEA.HANDLE_PUSH_TRIGGER", repo_trigger) + get_settings().set("GITEA.PUSH_COMMANDS", list(repo_commands)) + + def should_process_pr_logic(_body): + calls.append("filter") + return True + + @asynccontextmanager + async def record_slot(_key, **_kwargs): + calls.append("slot") + yield True + + monkeypatch.setattr(gitea_app, "apply_repo_settings", apply_repo_settings) + monkeypatch.setattr(gitea_app, "should_process_pr_logic", should_process_pr_logic) + monkeypatch.setattr(gitea_app, "push_trigger_slot", record_slot) + api_url = "https://gitea.example.com/org/repo/pulls/1" + + try: + await gitea_app.handle_pr_event( + {"pull_request": {"url": api_url}}, "pull_request", "synchronized", RecordingAgent() + ) + finally: + settings.set("GITEA", original_gitea) + + assert calls == ["settings", "filter"] + + +@pytest.mark.asyncio +async def test_gitea_repo_filter_rejects_push_before_reserving_slot(monkeypatch): + settings = get_settings() + original_gitea = copy.deepcopy(settings.get("GITEA")) + settings.set("GITEA.HANDLE_PUSH_TRIGGER", True) + settings.set("GITEA.PUSH_COMMANDS", ["/review"]) + calls = [] + + def apply_repo_settings(_url): + calls.append("settings") + + def should_process_pr_logic(_body): + calls.append("filter") + return False + + @asynccontextmanager + async def record_slot(_key, **_kwargs): + calls.append("slot") + yield True + + monkeypatch.setattr(gitea_app, "apply_repo_settings", apply_repo_settings) + monkeypatch.setattr(gitea_app, "should_process_pr_logic", should_process_pr_logic) + monkeypatch.setattr(gitea_app, "push_trigger_slot", record_slot) + api_url = "https://gitea.example.com/org/repo/pulls/1" + + try: + await gitea_app.handle_pr_event( + {"pull_request": {"url": api_url}}, "pull_request", "synchronized", RecordingAgent() + ) + finally: + settings.set("GITEA", original_gitea) + + assert calls == ["settings", "filter"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["opened", "reopened"]) +async def test_gitea_pr_event_applies_repo_settings_once_before_dispatch(monkeypatch, action): + settings = get_settings() + original_gitea = copy.deepcopy(settings.get("GITEA")) + original_is_auto_command = settings.get("CONFIG.IS_AUTO_COMMAND") + calls = [] + + def apply_repo_settings(_url): + calls.append("settings") + get_settings().set("GITEA.PR_COMMANDS", ["/review"]) + + def should_process_pr_logic(_body): + calls.append("filter") + return True + + class Agent: + async def handle_request(self, _url, command): + calls.append(command) + + monkeypatch.setattr(gitea_app, "apply_repo_settings", apply_repo_settings) + monkeypatch.setattr(gitea_app, "should_process_pr_logic", should_process_pr_logic) + monkeypatch.setattr(gitea_app, "prepare_command", lambda command: command) + api_url = "https://gitea.example.com/org/repo/pulls/1" + + try: + await gitea_app.handle_pr_event( + {"pull_request": {"url": api_url}}, "pull_request", action, Agent() + ) + finally: + settings.set("GITEA", original_gitea) + settings.set("CONFIG.IS_AUTO_COMMAND", original_is_auto_command) + + assert calls == ["settings", "filter", "/review"] + + @pytest.mark.asyncio @pytest.mark.parametrize( ("event", "draft", "feedback_on_draft_pr", "expected_commands"), diff --git a/uv.lock b/uv.lock index c48de07cc6..a4989b8854 100644 --- a/uv.lock +++ b/uv.lock @@ -330,11 +330,11 @@ wheels = [ [[package]] name = "certifi" -version = "2024.8.30" +version = "2026.7.22" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b0/ee/9b19140fe824b367c04c5e1b369942dd754c4c5462d5674002f75c4dedc1/certifi-2024.8.30.tar.gz", hash = "sha256:bec941d2aa8195e248a60b31ff9f0558284cf01a52591ceda73ea9afffd69fd9", size = 168507, upload-time = "2024-08-30T01:55:04.365Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/12/90/3c9ff0512038035f59d279fddeb79f5f1eccd8859f06d6163c58798b9487/certifi-2024.8.30-py3-none-any.whl", hash = "sha256:922820b53db7a7257ffbda3f597266d435245903d80737e34f8a45ff3e3230d8", size = 167321, upload-time = "2024-08-30T01:55:02.591Z" }, + { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, ] [[package]] @@ -405,63 +405,63 @@ wheels = [ [[package]] name = "charset-normalizer" -version = "3.4.8" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/20/56/10a88e00039537d74bd420f0457c52ab8f58a1af56126e3b9f1b1c8c4724/charset_normalizer-3.4.8.tar.gz", hash = "sha256:d9bf144d6faf12c70d58e47f7512992ae2882b820031d6cef68152deb645bf2d", size = 151790, upload-time = "2026-07-06T15:27:58.477Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/59/c2/39de60ef5687662f467bed3d1e6944c67a4f0d057141d0404002b8f405ae/charset_normalizer-3.4.8-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:faac37c4904598daa00cb4c9b32f3b4cc814fb5f145d7a531ceb4a70f2114132", size = 319040, upload-time = "2026-07-06T15:26:15.854Z" }, - { url = "https://files.pythonhosted.org/packages/a7/57/a9474c3aeaa337c8a330c0dc5df266527d56da3b189c029529f6b08af2a4/charset_normalizer-3.4.8-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f191c19a32dc6cec0fb8079789d786254653a9ce906fcab04ccd2eed07bba233", size = 215541, upload-time = "2026-07-06T15:26:17.265Z" }, - { url = "https://files.pythonhosted.org/packages/13/a9/be1ff7e81f6e086dced2a7a7a28b789be351d9796084ccaf6136a4ffafb3/charset_normalizer-3.4.8-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:05811b76943d477bb90822dedb5c4565cef70148847a59d574e2b35043aeb563", size = 236913, upload-time = "2026-07-06T15:26:18.482Z" }, - { url = "https://files.pythonhosted.org/packages/8f/75/d8c5eae93da26d463f9ebe46a4937ca44434dc2937a565b92437befb3d94/charset_normalizer-3.4.8-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3868a3e4ec1e40b419e060d063f93eac6f046fa21426c4816421223ae7dc8ab8", size = 232815, upload-time = "2026-07-06T15:26:19.734Z" }, - { url = "https://files.pythonhosted.org/packages/27/0d/98e301ca944bcca5e6bc312406b579c8a6d81546c1b494afb3a9478495d6/charset_normalizer-3.4.8-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:25f93d194eb6264c64416cabff46a91f6d99b97e7525a1b4f35c77a99e75cc68", size = 223995, upload-time = "2026-07-06T15:26:20.931Z" }, - { url = "https://files.pythonhosted.org/packages/fa/df/f5222366b76dcb31453a9bd922610c893540d0e729fd390439b0d3e972ee/charset_normalizer-3.4.8-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:ee6a62492f18d432cca031fabd158f400a8c25bf7b9458f50953393a2a23d97a", size = 208522, upload-time = "2026-07-06T15:26:22.214Z" }, - { url = "https://files.pythonhosted.org/packages/74/52/293220d59d8ddfb8aa56836b33bd6df58e70795d8a102a858c2984480f00/charset_normalizer-3.4.8-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:1c16cb4fc35e4b064f5ee78d849f15a550ada1729c3372916672e38f1f01d1d4", size = 219660, upload-time = "2026-07-06T15:26:23.493Z" }, - { url = "https://files.pythonhosted.org/packages/c3/2c/81a298e66f3d01e61bfc6f7064bbb553b067a9f1d979e5962bf00733069b/charset_normalizer-3.4.8-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2fbd0edb0426ab28e70fac9d1d4ef549eb5a64a2521f0428c441d75e4387e6c2", size = 218230, upload-time = "2026-07-06T15:26:24.629Z" }, - { url = "https://files.pythonhosted.org/packages/0d/45/f1dd2328cbc3340705f82072c09bd4c68d6e079191cde05810c1eac77eee/charset_normalizer-3.4.8-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:ccb9052771216170015f810b88065fd9e13b1e0b391f92abb9b47e0919a42aad", size = 210006, upload-time = "2026-07-06T15:26:25.837Z" }, - { url = "https://files.pythonhosted.org/packages/38/63/28697000620e117eb413424caaf60b6f98ddb1b09b2c11f7c0038d9936a7/charset_normalizer-3.4.8-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3809ba5d3cd02aca0894597f2669a825bdfe2229061515c128b0f4e5533b4ab5", size = 225771, upload-time = "2026-07-06T15:26:27.079Z" }, - { url = "https://files.pythonhosted.org/packages/12/a7/d5844e315f5f35e7938c415f07a1df144eed1cf993f1b43cc16c980c5b46/charset_normalizer-3.4.8-cp312-cp312-win32.whl", hash = "sha256:de63c31666a049f653ada24e800192e3c019e96bc7d70fb449a000bccf26a36f", size = 150922, upload-time = "2026-07-06T15:26:28.514Z" }, - { url = "https://files.pythonhosted.org/packages/9b/82/eb8b72f184b1e4986dd9daec15d7f6d9285a6728d2b07b7f04656829f473/charset_normalizer-3.4.8-cp312-cp312-win_amd64.whl", hash = "sha256:14a4bbe066f3fb05c6ba70e9cf9d34614b57a2fd70ea8c27cc30f34155e16a58", size = 162294, upload-time = "2026-07-06T15:26:29.745Z" }, - { url = "https://files.pythonhosted.org/packages/09/5a/ab810134aa41034a08ffe94c058102016e6ad9bce62f3cdba547b4723385/charset_normalizer-3.4.8-cp312-cp312-win_arm64.whl", hash = "sha256:2b5b0c0dca0a02c3f816f89abf18af3d20416dedbc3d3aa5f3981045f88ae7b0", size = 152409, upload-time = "2026-07-06T15:26:31.034Z" }, - { url = "https://files.pythonhosted.org/packages/11/49/fe5a8572a70cd9cba79f80af9388ac8c5c914ed4459b956f940244e499a5/charset_normalizer-3.4.8-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:057f8609f7341618c98e5aa9a6109fa116acff2a658497d47ab3325b5e8f2b08", size = 317424, upload-time = "2026-07-06T15:26:32.23Z" }, - { url = "https://files.pythonhosted.org/packages/e5/59/d71c96616b6825425a876f79f38fa440db30b32cc1166179a839f6259150/charset_normalizer-3.4.8-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c591f9a82adc5b89a039b90df74e43de2b9177fb46771172bed7b80722a70db0", size = 214723, upload-time = "2026-07-06T15:26:33.635Z" }, - { url = "https://files.pythonhosted.org/packages/c6/35/dc9eeb297f19b7b6ada39709ccb74937e6c51f0947958ae601a977cedd5d/charset_normalizer-3.4.8-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5b56f449132d9adefe55b87635d05177a914ed5d070438a74725e1d77a280002", size = 236200, upload-time = "2026-07-06T15:26:34.99Z" }, - { url = "https://files.pythonhosted.org/packages/0a/37/6775fe852b4acad8bf7e0575fbe8aa9f41b546e33251acbded3c04a6b0d9/charset_normalizer-3.4.8-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:1004c3b5831a301dadfb9e916f38e78e2ff3e08db24a1ad7c354db8ee3dea9c3", size = 231740, upload-time = "2026-07-06T15:26:36.498Z" }, - { url = "https://files.pythonhosted.org/packages/e4/28/1bcc3f5f3bac81532384adcfcdd9362c7f46a188a19deacc1ddaf7bdaa00/charset_normalizer-3.4.8-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:60883e22821d17c9e5b4f3ca1ef8074f766e3db28791f851b665929c515635c0", size = 222888, upload-time = "2026-07-06T15:26:38.161Z" }, - { url = "https://files.pythonhosted.org/packages/47/81/9f3993ca62ef090c58059da641e49e3129e74700a6a3beb58436cdb8d4b9/charset_normalizer-3.4.8-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:21e4dbb942c8a6342e2685f232dd2a7bc73465697bd26ead4f118271d28be383", size = 207649, upload-time = "2026-07-06T15:26:39.628Z" }, - { url = "https://files.pythonhosted.org/packages/11/f4/679f636bcbdc2d53d06b1f4039be310450dca95a9f76bbf22f09985556e8/charset_normalizer-3.4.8-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:84bcae14c65e645ca66b661339183d32b8c846a17c96e3e81ab3d346e1c498d4", size = 218908, upload-time = "2026-07-06T15:26:40.911Z" }, - { url = "https://files.pythonhosted.org/packages/7e/44/96e8c81867ba8a45ff893c8e7474c2d6b9633f7aa663da7901d040214d3e/charset_normalizer-3.4.8-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:b4e4d44b8287aa13a25e16e29393d494b0643b24894f7c8266c6f6788dd36337", size = 217096, upload-time = "2026-07-06T15:26:42.148Z" }, - { url = "https://files.pythonhosted.org/packages/a5/bf/4d53f04f29bdb22601701f4f9f4d038edfb27976c296fcb7400c02736a6e/charset_normalizer-3.4.8-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:f68545d1b267dbfafd5d253b6d1cb161562c4e61ab25b5c4cdb7d9e5923e441e", size = 209355, upload-time = "2026-07-06T15:26:43.557Z" }, - { url = "https://files.pythonhosted.org/packages/ad/60/92b3f630798d777fa880ad289a3f9f2fc663e4b4beb24783c53318820254/charset_normalizer-3.4.8-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1936e48214adea74922a20c8ab41b1393ae27cc9e329eb1f0b937d3416824f36", size = 224732, upload-time = "2026-07-06T15:26:44.892Z" }, - { url = "https://files.pythonhosted.org/packages/bd/85/eafa0a3c7bb6fe9f02f4c7901f02071933cac85ee634197e17280818c6de/charset_normalizer-3.4.8-cp313-cp313-win32.whl", hash = "sha256:1f8e3521860187d597f3867d8466da225b9179ea2833bb26de1bb026144d07c3", size = 150358, upload-time = "2026-07-06T15:26:46.153Z" }, - { url = "https://files.pythonhosted.org/packages/f5/8c/879fafff7b47bb1166d289f2d2472cb31b9922f9f4ca1f392edf85ec16be/charset_normalizer-3.4.8-cp313-cp313-win_amd64.whl", hash = "sha256:8b654b6f52a0a9a6be38e88f3e1dc68f1093ebeb2abbadafc7c82da0786a34be", size = 161685, upload-time = "2026-07-06T15:26:47.423Z" }, - { url = "https://files.pythonhosted.org/packages/c3/46/b57c7e778a7b578f28d35fd38544687d4f8d9c019585eebc5ad936073fad/charset_normalizer-3.4.8-cp313-cp313-win_arm64.whl", hash = "sha256:d2d5a250ee26e29468b7607d97479221b069fa8aaf6f929ac84ec0e962e15154", size = 152333, upload-time = "2026-07-06T15:26:48.689Z" }, - { url = "https://files.pythonhosted.org/packages/1c/bc/0a8540b8cd494951cca1428606373942803f5ffcec40fe798f819c5a8adb/charset_normalizer-3.4.8-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:77e993ecf65f21ab1f82266ff5e84a7de2c879e7d9b8bc006009df83f22a1d5e", size = 316993, upload-time = "2026-07-06T15:26:49.962Z" }, - { url = "https://files.pythonhosted.org/packages/0e/99/a0868f0a1f0a045fd374d1f2cf7042d8ad5d7fb4dd1f4ac7365e319f7e32/charset_normalizer-3.4.8-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:524939917f17f6de502dfda30b472550965740d7f126659d4c4f8dd1569cce22", size = 215638, upload-time = "2026-07-06T15:26:51.338Z" }, - { url = "https://files.pythonhosted.org/packages/e8/e9/43c4d09a09b5557cc5fe1d87c9d96f86a3942aec0517d2b5408cef87ca75/charset_normalizer-3.4.8-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a4508989ba8e2ce43ef989453d18188b261546e8188cbdd4ef451fb9e4c3b467", size = 236456, upload-time = "2026-07-06T15:26:52.531Z" }, - { url = "https://files.pythonhosted.org/packages/e2/67/492ca98b3ab785b736b5da10c1bc233e1c8fec6c0cdb29b482c38bfc52a2/charset_normalizer-3.4.8-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:9e44127f7d11eee4548ad2cdf1f4e1b6eaaddd5cb92d15ad65f6ecc9bcf403ab", size = 232253, upload-time = "2026-07-06T15:26:53.838Z" }, - { url = "https://files.pythonhosted.org/packages/2d/fd/1e6eff58c14f1aace1e26d80defbeaea2d35e075dbe4b611111ee4b47fa8/charset_normalizer-3.4.8-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bb90317359f7e67bb6df615999a95e0980877468e617ddce8b6c2f8e7fe60d95", size = 222886, upload-time = "2026-07-06T15:26:55.009Z" }, - { url = "https://files.pythonhosted.org/packages/40/7a/90056a5326b0c4b9a3f924d337729c344c11542e5bc7191e50410db61587/charset_normalizer-3.4.8-cp314-cp314-manylinux_2_31_armv7l.whl", hash = "sha256:35d9e7a9c960520ae89d1f4e305d1c047a74dea2e0f73a0e84f879356c2e8776", size = 206482, upload-time = "2026-07-06T15:26:56.306Z" }, - { url = "https://files.pythonhosted.org/packages/18/ff/94761d31a33878dbb5008ddbd918615061fcf5c0a612aa3075450e60f628/charset_normalizer-3.4.8-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:92e322b054c7ff886f78feab7360736bb45de2e18cf4a0ee84e8fc5a08d53a19", size = 218929, upload-time = "2026-07-06T15:26:57.422Z" }, - { url = "https://files.pythonhosted.org/packages/2b/dc/00b9675acd7c4b926b9102ee3f0d1a570ce943901be73b87485001393fe1/charset_normalizer-3.4.8-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:3c0086d97094363556206dc3bcf43f7edcfc043ea7a568a46f45efea74858bd1", size = 218069, upload-time = "2026-07-06T15:26:58.719Z" }, - { url = "https://files.pythonhosted.org/packages/04/11/94ada5a0482ee4bf688d04be4c7d6fd945d37370d04a95671040dfe2b416/charset_normalizer-3.4.8-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:0752c849b51198267df2aba013c4de3a2955bd014a4fd70828809946c1acbc0c", size = 207146, upload-time = "2026-07-06T15:27:00.058Z" }, - { url = "https://files.pythonhosted.org/packages/7c/f7/246bd36762207ab4752cd436b64e5d81a1668b15ddea7b5b2d0e8545e727/charset_normalizer-3.4.8-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2a4707e09eca11e81ece4fced600c5a0a801f568b962244f6f517bc274745fc9", size = 224896, upload-time = "2026-07-06T15:27:01.599Z" }, - { url = "https://files.pythonhosted.org/packages/6f/f7/3510622d1fbe13b0ebf827c475e40a27e2be427140d792878b63ab6425cc/charset_normalizer-3.4.8-cp314-cp314-win32.whl", hash = "sha256:8ea67f427c073ae3da0923aa55f3715131fa613a61a7f2f8d762bde75eaf00ae", size = 150851, upload-time = "2026-07-06T15:27:02.964Z" }, - { url = "https://files.pythonhosted.org/packages/32/2b/9ce65dd21672b55cf800cca5f4433afa1586fda1d78731067ec9ec544c62/charset_normalizer-3.4.8-cp314-cp314-win_amd64.whl", hash = "sha256:ff71018850863362e5c7533769d0a9f77715c31af1502d523630ce822922f5c9", size = 162549, upload-time = "2026-07-06T15:27:04.249Z" }, - { url = "https://files.pythonhosted.org/packages/2f/34/9a5967eed666a88f31a0866884606d9ec3c2cd6091e2ccd7e0b4c4176c35/charset_normalizer-3.4.8-cp314-cp314-win_arm64.whl", hash = "sha256:44464e66f4da2f21dea7145c7693f9f60717ca4794a954dea5bf8c2c932678bd", size = 153079, upload-time = "2026-07-06T15:27:05.608Z" }, - { url = "https://files.pythonhosted.org/packages/02/4f/aa44cc81d8987f105352c74c0bf919007f8b80e9880d28bcf0393c1a816e/charset_normalizer-3.4.8-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:50a0c2e58ad2c203adb616fef28941b7e13716adbc25e0dfaeec29f5afe6382f", size = 338586, upload-time = "2026-07-06T15:27:06.86Z" }, - { url = "https://files.pythonhosted.org/packages/1d/2b/b0392e2b235c08ff0623d905c2ee8ac820620544043c1ce92ce0b3d64c55/charset_normalizer-3.4.8-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a1e589fdb95c76f08288bbb346230cdd8994db74903db6637b380f7b5fc9336", size = 222764, upload-time = "2026-07-06T15:27:08.23Z" }, - { url = "https://files.pythonhosted.org/packages/a2/a1/7d466879190731f5559662c22232646f2ae2dace2323c3e5aefcf78d458a/charset_normalizer-3.4.8-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b3d7c887444c5a7ef0d68d358d81e758a850bc626f8e639e2ca5667153272b20", size = 241331, upload-time = "2026-07-06T15:27:09.512Z" }, - { url = "https://files.pythonhosted.org/packages/70/17/8b89e797137aa28c8fb0bafbafc243246a7afe21620a13b00e37624ece1d/charset_normalizer-3.4.8-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:65c389b96c0cfff3a3f0458fa1c7ce554a30e23101a88a49f03997afce6a929f", size = 239323, upload-time = "2026-07-06T15:27:10.86Z" }, - { url = "https://files.pythonhosted.org/packages/d7/98/1c1940730ed22d50983be4e243c722c89d5136d6f073bd840d1128bfddcb/charset_normalizer-3.4.8-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:593403fc47dcdf55e2987b2e3cc2e064127e2b908929f1f18b2e4a4652cbd780", size = 229964, upload-time = "2026-07-06T15:27:12.113Z" }, - { url = "https://files.pythonhosted.org/packages/53/2d/bb8e81b7ff603d3f77e9a8a5d1ad34fcabbf3c54d300c29d99fba581fa23/charset_normalizer-3.4.8-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:606088e9fa2b7469ab9c42d4da8e05a415622a07714edd2fcd8fed48dda4c853", size = 212405, upload-time = "2026-07-06T15:27:13.447Z" }, - { url = "https://files.pythonhosted.org/packages/ce/1f/e52a3a53b13da591bb8f21d29e63877268eadf20686b7762351d4b89062c/charset_normalizer-3.4.8-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0317406326fed512f42a1632ad91a96228a7616c06547666a6dd79967f1bd6ca", size = 226918, upload-time = "2026-07-06T15:27:14.89Z" }, - { url = "https://files.pythonhosted.org/packages/e8/f9/32996d79c57189af9722fe618f46d8a86b7be035ca98887b8d0c3821f141/charset_normalizer-3.4.8-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:b67d50ee47e5c57a0064a9cb575b963a7125819dfd1fd094d44d378fff94659b", size = 225113, upload-time = "2026-07-06T15:27:16.125Z" }, - { url = "https://files.pythonhosted.org/packages/d6/d2/9248c18e695696513774523a794cfb8b677521ce9ad7554d301cb10a9b20/charset_normalizer-3.4.8-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:79e402b869f270140afa5e2b0e2ac100585358d812fe3dd093d424f7a72964e0", size = 214966, upload-time = "2026-07-06T15:27:17.418Z" }, - { url = "https://files.pythonhosted.org/packages/1e/9d/4b19432d406179a40f924691906ee5b15ac664b408971c973295192444ea/charset_normalizer-3.4.8-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:2970b9f7ab69ec3a0423ec6b6ac718e79fbf4a282c0bc103ef88c1ef50dfa15a", size = 231699, upload-time = "2026-07-06T15:27:19.131Z" }, - { url = "https://files.pythonhosted.org/packages/be/41/bdbdf71e8c3ccff10ef3cc2bb9467a7fdb3dc94b9a406d1a3c44afd39632/charset_normalizer-3.4.8-cp314-cp314t-win32.whl", hash = "sha256:458c2972a78043b7261c9726670029f15f722e70669bcbe961153a01968f589f", size = 155333, upload-time = "2026-07-06T15:27:20.681Z" }, - { url = "https://files.pythonhosted.org/packages/bd/f8/e05c69323bd50091ec39f5f885385b884624b0131a6885a0c83a6217ba7a/charset_normalizer-3.4.8-cp314-cp314t-win_amd64.whl", hash = "sha256:0c926329a1df7cd56d7d8349fe354460d20aefd2e394c9e159e479d018b2b359", size = 167378, upload-time = "2026-07-06T15:27:22.042Z" }, - { url = "https://files.pythonhosted.org/packages/c2/04/cbaf1a2f5e2bbf70760e774380cbf052b10849fc35e770905df31af5cf00/charset_normalizer-3.4.8-cp314-cp314t-win_arm64.whl", hash = "sha256:2232baea80a2b01783679fed4e625ccdb19a974f44c9cf0fba21a777a4c8179c", size = 157782, upload-time = "2026-07-06T15:27:23.312Z" }, - { url = "https://files.pythonhosted.org/packages/23/52/d5bee5b6ea81882d549b566d2545b044bbcbc33fe5fbe001008a7e745a21/charset_normalizer-3.4.8-py3-none-any.whl", hash = "sha256:b7c1fb310df524e01fbe84d43b7f95aa4f808f8eaa0dafc185f64ba395e37d54", size = 64279, upload-time = "2026-07-06T15:27:57.043Z" }, +version = "3.4.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bd/2a/23f34ec9d04624958e137efdc394888716353190e75f25dd22c7a2c7a8aa/charset_normalizer-3.4.9.tar.gz", hash = "sha256:673611bbd43f0810bec0b0f028ddeaaa501190339cac411f347ac76917c3ae7b", size = 152439, upload-time = "2026-07-07T14:34:58.454Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/70/4a/ecbd131485c07fcdfad54e28946d513e3da22ef3b4bd854dcafae54ec739/charset_normalizer-3.4.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:45b0cc4e3556cd875e09102988d1ab8356c998b596c9fced84547c8138b487a0", size = 319300, upload-time = "2026-07-07T14:33:15.666Z" }, + { url = "https://files.pythonhosted.org/packages/ec/96/5d9364e3342d69f3a045e1777bc47c85c383e6e9466d561b33fdb419d1f9/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9b2aff1c7b3884512b9512c3eaadd9bab39fb45042ffaaa1dd08ff2b9f8109d9", size = 215802, upload-time = "2026-07-07T14:33:17.031Z" }, + { url = "https://files.pythonhosted.org/packages/4b/4c/5361f9aa7f2cb58d94f2ab831b3d493f69efb1d239654b4744e3c09527cb/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9104ed0bd76a429d46f9ec0dbc9b08ad1d2dcdf2b00a5a0daa1c145329b35b44", size = 237171, upload-time = "2026-07-07T14:33:18.576Z" }, + { url = "https://files.pythonhosted.org/packages/50/78/ce342ca4ff30b2eb49fe6d9578df85974f90c67d294113e94efdd9664cbd/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7b86a2b16095d250c6f58b3d9b2eee6f4147754344f3dab0922f7c9bf7d226c9", size = 233075, upload-time = "2026-07-07T14:33:20.084Z" }, + { url = "https://files.pythonhosted.org/packages/01/c4/4fa4c8b3097a11f3c5f09a35b72ed6855fb1d332469504962ab7bafcc702/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5e226f6218febc71f6c1fc2fafb91c226f75bdc1d8fb12d66823716e891608fd", size = 224256, upload-time = "2026-07-07T14:33:21.747Z" }, + { url = "https://files.pythonhosted.org/packages/87/3a/ad914516df7e358a81aae018caa5e0470ba827fa6d763b1d2e87d920a5f6/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:90c44bc373b7687f6948b693cceaea1348ae0975d7474746559494468e3c1d84", size = 208784, upload-time = "2026-07-07T14:33:23.313Z" }, + { url = "https://files.pythonhosted.org/packages/d7/74/3c12f9755717dfe5c5c87da63f35d765fa0c00382ec26bf23f7fae34f2ba/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9cdef90ae47919cae358d8ab15797a800ed41da7aba5d72419fb510729e2ed4b", size = 219928, upload-time = "2026-07-07T14:33:24.814Z" }, + { url = "https://files.pythonhosted.org/packages/33/9a/895095b83e7907abd6d3d99aad3a38ad0d9686cc186cb0c94c24320fe63e/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:60f44ade2cf573dad7a277e6f8ca9a51a21dda572b13bd7d8539bb3cd5dbedde", size = 218489, upload-time = "2026-07-07T14:33:26.42Z" }, + { url = "https://files.pythonhosted.org/packages/a1/34/ef5c05f412f42520d7709b7d3784d19640839eb7366ded1755511585429f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:a1786910334ed46ab1dd73222f2cd1e05c2c3bb39f6dddb4f8b36fc382058a39", size = 210267, upload-time = "2026-07-07T14:33:27.952Z" }, + { url = "https://files.pythonhosted.org/packages/83/dc/9b29fa4412b318bf3bfea985c35d67eb55e04b59a7c3f2237168b0e0be6f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:03d07803992c6c7bbc976327f34b18b6160327fc81cb82c9d504720ac0be3b62", size = 226030, upload-time = "2026-07-07T14:33:29.397Z" }, + { url = "https://files.pythonhosted.org/packages/0e/42/6dbc00b8cd16011691203e33570fa42ed5746599a2e878112d16eab403a3/charset_normalizer-3.4.9-cp312-cp312-win32.whl", hash = "sha256:78841cccf1af7b40f6f716338d50c0902dbe88d9f800b3c973b7a9a0a693a642", size = 151185, upload-time = "2026-07-07T14:33:30.781Z" }, + { url = "https://files.pythonhosted.org/packages/80/cc/f920afd1a23c58ccd53c1d36085a71893a4737ff5e66e0371efab6809850/charset_normalizer-3.4.9-cp312-cp312-win_amd64.whl", hash = "sha256:4b3dac63058cc36820b0dd072f89898604e2d39686fe05321729d00d8ac185a0", size = 162557, upload-time = "2026-07-07T14:33:32.176Z" }, + { url = "https://files.pythonhosted.org/packages/f0/e6/0386d43a261ff4e4b30c5857af7df877254b46bec7b9d1b74b6bf969a90b/charset_normalizer-3.4.9-cp312-cp312-win_arm64.whl", hash = "sha256:78fa18e436a1a0e58dbd7e02fc4473f3f32cceb12df9dfca542d075961c307d2", size = 152665, upload-time = "2026-07-07T14:33:33.711Z" }, + { url = "https://files.pythonhosted.org/packages/b2/06/97ec2aeae780b31d742b6352218b43841a6871e2564578ca522dce4a45c3/charset_normalizer-3.4.9-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:440eede837960000d74978f0eba527be106b5b9aee0daf779d395276ed0b0614", size = 317688, upload-time = "2026-07-07T14:33:35.408Z" }, + { url = "https://files.pythonhosted.org/packages/d0/39/8ff066c672434225f8d25f8b739f992af250944392173dcc88362681c9bf/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:21e764fd1e70b6a3e205a0e46f3051701f98a8cb3fad66eeb80e48bb502f8698", size = 214982, upload-time = "2026-07-07T14:33:36.996Z" }, + { url = "https://files.pythonhosted.org/packages/92/8f/3a47a3667c83c2df9483d91644c6c107de3bf8874aa1793da9d3012eb986/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e4fd89cc178bced6ad29cb3e6dd4aa63fa5017c3524dbd0b25998fb64a87cc8b", size = 236460, upload-time = "2026-07-07T14:33:38.536Z" }, + { url = "https://files.pythonhosted.org/packages/f1/60/b22cdbee7e4013dab8b0d7647fc6181120fbbbc8f7025c226d15bd5a47fc/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bd47ba7fc3ca94896759ea0109775132d3e7ab921fbf54038e1bab2e46c313c9", size = 232003, upload-time = "2026-07-07T14:33:40.059Z" }, + { url = "https://files.pythonhosted.org/packages/ea/f8/72eb13dcabe7257035cea8aefd922caad2f110d252bf9f67c4c2ca763aee/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:84fd18bcc17526fc2b3c1af7d2b9217d32c9c04448c16ec693b9b4f1985c3d33", size = 223149, upload-time = "2026-07-07T14:33:41.631Z" }, + { url = "https://files.pythonhosted.org/packages/b0/3e/faee8f9de92b14ee1198e9163252bb15efee7301b31256a3b6d9ebfdd0dd/charset_normalizer-3.4.9-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:5b10cd92fc5c498b35a8635df6d5a100207f88b63a4dc1de7ef9a548e1e2cd63", size = 207901, upload-time = "2026-07-07T14:33:43.209Z" }, + { url = "https://files.pythonhosted.org/packages/3a/25/45f30093ae27dd7b92a793b61882a38685f993700113ca36e0c9c14965e1/charset_normalizer-3.4.9-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a4fbdde9dd4a9ce5fd52c2b3a347bb50cc89483ef783f1cb00d408c13f7a96c0", size = 219176, upload-time = "2026-07-07T14:33:44.725Z" }, + { url = "https://files.pythonhosted.org/packages/48/18/c8f397329c35e32f6a837e488986f4ae03bd2abebc453b48714991630c2f/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:416c229f77e5ea25b3dfd4b582f8d73d7e43c22320302b9ab128a2d3a0b38efe", size = 217356, upload-time = "2026-07-07T14:33:46.192Z" }, + { url = "https://files.pythonhosted.org/packages/86/7e/5ce0bba863470fd1902d5e5843968951bddf38abe4742fc97116ef4598b3/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:75286256590a6320cf106a0d28970d3560aad9ee09aa7b34fb40524792436d35", size = 209614, upload-time = "2026-07-07T14:33:47.705Z" }, + { url = "https://files.pythonhosted.org/packages/6c/ef/2473d3c4d869155be4af1191111d59c4d5c4e0173026f7e85b176e23bf65/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:69b157c5d3292bcd443faca052f3096f637f1e074b98212a933c074ae23dc3b8", size = 224991, upload-time = "2026-07-07T14:33:49.238Z" }, + { url = "https://files.pythonhosted.org/packages/d0/a3/53ddae3db108a088156aa8ddfafd411ebbc1340f48c5573f697b27f69a39/charset_normalizer-3.4.9-cp313-cp313-win32.whl", hash = "sha256:51307f5c71007673a2bf8232ad973483d281e74cb99c8c5a990af1eefa6277d9", size = 150622, upload-time = "2026-07-07T14:33:50.711Z" }, + { url = "https://files.pythonhosted.org/packages/e8/ef/6953a77c7cf2c2ff9998e6f575ab3e380119f100223381565a4f94c1f836/charset_normalizer-3.4.9-cp313-cp313-win_amd64.whl", hash = "sha256:fe2c7201c642b7c308f1675355ad7ff7b66acfe3541625efe5a3ad38f29d6115", size = 161947, upload-time = "2026-07-07T14:33:52.197Z" }, + { url = "https://files.pythonhosted.org/packages/6e/fb/d560d1d1555debbfe7849d9cac6145c1b537709d79576bf22557ed803b82/charset_normalizer-3.4.9-cp313-cp313-win_arm64.whl", hash = "sha256:611057cc5d5c0afc743ba8be6bd828c17e0aaa8643f9d0a9b9bb7dea80eb8012", size = 152594, upload-time = "2026-07-07T14:33:53.486Z" }, + { url = "https://files.pythonhosted.org/packages/7e/8d/496817fa0944239ecae662dd57ea765cfeaec6a735f9f025d4b7b72e7143/charset_normalizer-3.4.9-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:0327fcd59a935777d83410750c50600ee9571af2846f71ce40f25b13da1ef380", size = 317253, upload-time = "2026-07-07T14:33:54.994Z" }, + { url = "https://files.pythonhosted.org/packages/2b/f9/ef4a69ea338ad3c0deceea0f5f7d2380ae8b52132b06d652cb0d2cd86706/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a79d9f4d8001473a30c163556b3c3bfebec837495a412dde78b51672f6134f9", size = 215898, upload-time = "2026-07-07T14:33:56.334Z" }, + { url = "https://files.pythonhosted.org/packages/8c/e7/5ddfd76fc061eb52de219658a4aa431cbacadf0a0219c8854f00da50d289/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:33bdcc2a32c0a0e861f60841a512c8acc658c87c2ac59d89e3a46dacf7d866e4", size = 236718, upload-time = "2026-07-07T14:33:57.9Z" }, + { url = "https://files.pythonhosted.org/packages/49/ba/768fa3f36048d81c477a0ce61f813bc1454d80917ccfe550abd9f44f5e24/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f840ed6d8ecba8255df8c42b87fadeda98ddfc6eeec05e2dc66e26d46dd6f58a", size = 232519, upload-time = "2026-07-07T14:33:59.811Z" }, + { url = "https://files.pythonhosted.org/packages/f4/c4/b3e049d2aa3766180c78507110543d9d50894cc97f57de543f1be521dcdc/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c25fe15c70c59eb7c5ce8c06a1f3fa1da0ecc5ea1e7a5922c40fd2fa9b0d5046", size = 223143, upload-time = "2026-07-07T14:34:01.517Z" }, + { url = "https://files.pythonhosted.org/packages/19/79/55c32d06d76ae4feafe053f061f3e3ab70bcf19f4007797ce8c3efda7830/charset_normalizer-3.4.9-cp314-cp314-manylinux_2_31_armv7l.whl", hash = "sha256:f7fb7d750cfa0a070d2c24e831fd3481019a60dd317ea2b39acbcebc08b6ed81", size = 206742, upload-time = "2026-07-07T14:34:03.04Z" }, + { url = "https://files.pythonhosted.org/packages/10/e0/47c079dd82d217c807479cd59ffd30af56307ea31c108b75758970459ad3/charset_normalizer-3.4.9-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4d1c96a7a18b9690a4d46df09e3e3382406ae3213727cd1019ebade1c4a81917", size = 219191, upload-time = "2026-07-07T14:34:04.657Z" }, + { url = "https://files.pythonhosted.org/packages/42/ab/b9bc2e77d6b44a7e46ef62ec5cac1c9a6ba7b9135a5d560f002696ec9995/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a4cfde78a9f2880208d16a93b795726a3017d5977e08d1e162a7a31322479c41", size = 218328, upload-time = "2026-07-07T14:34:06.115Z" }, + { url = "https://files.pythonhosted.org/packages/f1/78/c9c71d599f5aa2d42bcdd35cbbd46d7f535351a57e40ff7d8e5a7e219401/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:d4d6fcde76f94f5cb9e43e9e9a61f16dacefd228cbbf6f1a09bd9b219a92f1a1", size = 207406, upload-time = "2026-07-07T14:34:07.554Z" }, + { url = "https://files.pythonhosted.org/packages/f6/39/c914445c321a845097ce4f6ac7de9a18228a77b766272125a1ce00d851eb/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:898f0e9068ca27d37f8e83a5b962821df851532e6c4a7d615c1c033f9da6eedf", size = 225157, upload-time = "2026-07-07T14:34:09.061Z" }, + { url = "https://files.pythonhosted.org/packages/9b/f2/c0d4b8508565a36bc5c624e88ed297f5b0b1095011034d7f5b83a69908b5/charset_normalizer-3.4.9-cp314-cp314-win32.whl", hash = "sha256:c1c948747b03be832dceed96ca815cef7360de9aa19d37c730f8e3f6101aca48", size = 151095, upload-time = "2026-07-07T14:34:10.901Z" }, + { url = "https://files.pythonhosted.org/packages/49/fd/a1d26144398c67486422a72bf5812cda22cb4ccfcd95a290fb41ceb4b8e2/charset_normalizer-3.4.9-cp314-cp314-win_amd64.whl", hash = "sha256:16b65ea0f2465b6fb52aa22de5eca612aa964ddfec00a912e26f4656cbef890b", size = 162796, upload-time = "2026-07-07T14:34:12.47Z" }, + { url = "https://files.pythonhosted.org/packages/20/95/d75e82f8ce9fd323ebf059c16c9aadefb22a1ecde13b7840b35835e4886c/charset_normalizer-3.4.9-cp314-cp314-win_arm64.whl", hash = "sha256:40a126142a56b2dfc0aacbad1de8310cbf60da7656db0e6b16eebd48e3e93519", size = 153334, upload-time = "2026-07-07T14:34:14.044Z" }, + { url = "https://files.pythonhosted.org/packages/00/5e/17398df3a139985ba9d11ed072531986f408c8fca952835ef1ab1820c02b/charset_normalizer-3.4.9-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:609b3ba8fcc0fb5ab7af00719d0fb6ad0cb518e48e7712d12fd68f1327951198", size = 338848, upload-time = "2026-07-07T14:34:15.688Z" }, + { url = "https://files.pythonhosted.org/packages/cd/91/7253a32e86b7e1d1239b1b36ba6dd0f021a21107ab33054b53119cc083b9/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:51447e9aa2684679af07ca5021c3db526e0284347ebf4ffcec1154c3350cfe32", size = 223022, upload-time = "2026-07-07T14:34:17.248Z" }, + { url = "https://files.pythonhosted.org/packages/cb/32/2e64bd2be10e89c61e57ebe6a93fd98ae88eb7ebe414b5121f22c96c69eb/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cc1b0fff8ead343dae06305f954eb8468ba0ec1a97881f42489d198e4ce3c632", size = 241590, upload-time = "2026-07-07T14:34:18.813Z" }, + { url = "https://files.pythonhosted.org/packages/3d/ef/d96ec496cfea0c21db43b0ad03891308b02388d054cc902cf0e5a1ad6a88/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fa36ec09ef71d158186bc79e359ff5fdd6e7996fe8ab638f00d6b93139ba4fcf", size = 239584, upload-time = "2026-07-07T14:34:20.52Z" }, + { url = "https://files.pythonhosted.org/packages/d4/ce/9af95f7876194bd7a14e3dfe4a4de2e0bff02666a3910d72beafd06cc297/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:df115d4d83168fdf2cae48ef1ff6d1cb4c466364e30861b37121de0f3bf1b990", size = 230224, upload-time = "2026-07-07T14:34:22.189Z" }, + { url = "https://files.pythonhosted.org/packages/52/94/af74dde74a3996bd959c350709bfe50e297823d70a8c1cbd54b838880863/charset_normalizer-3.4.9-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:f86c6358749bd4fda175388691e3ba8c46e24c5347d0afd20f9b7edfc9faf07d", size = 212667, upload-time = "2026-07-07T14:34:23.857Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f0/f1c4fe746c395922961b5916ed1d7d6e7d4c84851d19ed43cc89980ec953/charset_normalizer-3.4.9-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:32286a2c8d167e897177b673176c1e3e00d4057caf5d2b64eef9a3666b03018e", size = 227179, upload-time = "2026-07-07T14:34:25.586Z" }, + { url = "https://files.pythonhosted.org/packages/e4/56/6c745619ac397e8871e2bcd3cea1eec86b877488f33888b3aef5c3ed506e/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:83aed2c10721ddd90f68140685391b50811a880af20654c59af6b6c66c40513c", size = 225372, upload-time = "2026-07-07T14:34:27.212Z" }, + { url = "https://files.pythonhosted.org/packages/78/ad/98aae8630ac71f16711968e38a5acfecce41b778bf2f0312851020f565a8/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:cd6c3d4b783c556fa00bf540854e42f135e2f256abd29669fcd0da0f2dec79c2", size = 215222, upload-time = "2026-07-07T14:34:28.774Z" }, + { url = "https://files.pythonhosted.org/packages/f7/40/9593d54209765207a7f11073c06494c1721e4ca4a0a426c597679bf7f91e/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ee2f2a527e3c1a6e6411eb4209642e138b544a2d72fe5d0d76daf77b24063534", size = 231958, upload-time = "2026-07-07T14:34:30.345Z" }, + { url = "https://files.pythonhosted.org/packages/b1/27/693ee5e8a18191eb38647360c51cd505013e2bd3b366aa43fd5344c21e3c/charset_normalizer-3.4.9-cp314-cp314t-win32.whl", hash = "sha256:0d861473f743244d349b50f850d10eb87aeb22bbdcc8e64f79273c94af5a8226", size = 155580, upload-time = "2026-07-07T14:34:31.884Z" }, + { url = "https://files.pythonhosted.org/packages/80/3f/bd97d3d9c613013d07cb7733d299385b41df37f0471310f5a73dc359f0b8/charset_normalizer-3.4.9-cp314-cp314t-win_amd64.whl", hash = "sha256:9b8e0f3107e2200b76f6054de99016eac3ee6762713587b36baaa7e4bd2ae177", size = 167620, upload-time = "2026-07-07T14:34:33.438Z" }, + { url = "https://files.pythonhosted.org/packages/3d/c6/eee9dca4439b1061f76373f06ea855678cc4a64c1c3c90b50e479edbb8eb/charset_normalizer-3.4.9-cp314-cp314t-win_arm64.whl", hash = "sha256:19ac87f93086ce37b86e098888555c4b4bc48102279bae3350098c0ed664b501", size = 158037, upload-time = "2026-07-07T14:34:35.018Z" }, + { url = "https://files.pythonhosted.org/packages/98/2b/f97f1c193fb855c345d678f5077d6926034db0722df74c8f057020e05a25/charset_normalizer-3.4.9-py3-none-any.whl", hash = "sha256:68e5f26a1ad57ded6d1cfb85331d1c1a195314756471d97758c48498bb4dcdf5", size = 64538, upload-time = "2026-07-07T14:34:56.993Z" }, ] [[package]] @@ -872,57 +872,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/22/4222d7ddf3da30f363edaa98e329c2bce6c65497c9cb2810931c8b2c0fbc/fsspec-2026.6.0-py3-none-any.whl", hash = "sha256:02e0b71817df9b2169dc30a16832045764def1191b43dcff5bb85bdee212d2a1", size = 203949, upload-time = "2026-06-16T01:57:26.358Z" }, ] -[[package]] -name = "gcsfs" -version = "2025.12.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version < '3.13' and sys_platform == 'win32'", - "python_full_version < '3.13' and sys_platform == 'emscripten'", - "python_full_version < '3.13' and sys_platform != 'emscripten' and sys_platform != 'win32'", -] -dependencies = [ - { name = "aiohttp" }, - { name = "decorator" }, - { name = "fsspec", version = "2025.12.0", source = { registry = "https://pypi.org/simple" } }, - { name = "google-auth" }, - { name = "google-auth-oauthlib" }, - { name = "google-cloud-storage", version = "2.10.0", source = { registry = "https://pypi.org/simple" } }, - { name = "google-cloud-storage-control" }, - { name = "requests" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/3e/3e/453b42bbcda2177daba97ebc5202faf5c57cc0c2bb4aac7081c12b0471da/gcsfs-2025.12.0.tar.gz", hash = "sha256:9a9c1c32b7899a3967ba30a7e9422cdfda596266f03476cdf5545402b8af4cd5", size = 95148, upload-time = "2025-12-03T15:44:59.703Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b1/bf/e664cbeab8f2a8b097c0604252760410fde191fe6ac7d5081b29e601ac52/gcsfs-2025.12.0-py3-none-any.whl", hash = "sha256:e06aaec53797dc6b83d5cc90c4d3ae7247b4ee0cf8d8b1ce50e8d6b78e3a9aea", size = 41204, upload-time = "2025-12-03T15:44:58.464Z" }, -] - -[[package]] -name = "gcsfs" -version = "2026.7.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.14' and sys_platform == 'win32'", - "python_full_version >= '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", -] -dependencies = [ - { name = "aiohttp" }, - { name = "decorator" }, - { name = "fsspec", version = "2026.6.0", source = { registry = "https://pypi.org/simple" } }, - { name = "google-auth" }, - { name = "google-auth-oauthlib" }, - { name = "google-cloud-storage", version = "3.12.0", source = { registry = "https://pypi.org/simple" } }, - { name = "google-cloud-storage-control" }, - { name = "requests" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d3/33/58e23c6f492e091c2f01ae1bffe331aa97e18a0b3d24ff90ef648335ddc8/gcsfs-2026.7.0.tar.gz", hash = "sha256:2df36ce1e9010e76dc4295774030495a97496838483c619e9f63bc0ab7bdc770", size = 1265874, upload-time = "2026-07-06T15:02:30.683Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ef/eb/a5c4b39903928e6be7735cc7aa31b6c3bbf3d887419e1e7ac751dba02233/gcsfs-2026.7.0-py3-none-any.whl", hash = "sha256:8ce6c08ea64302d8fd4bc23a615ccd9e7752541cde1a88ec8362109bd79cde78", size = 91131, upload-time = "2026-07-06T15:02:29.208Z" }, -] - [[package]] name = "gitdb" version = "4.0.12" @@ -1002,19 +951,6 @@ requests = [ { name = "requests" }, ] -[[package]] -name = "google-auth-oauthlib" -version = "1.4.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "google-auth" }, - { name = "requests-oauthlib" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/70/18/90c7fac516e63cf2058166fce0c88c353647c677b51cc036c09c49bb5cbb/google_auth_oauthlib-1.4.0.tar.gz", hash = "sha256:18b5e28880eb8eba9065c436becdc0ee8e4b59117a73a510679c82f70cd363d2", size = 21675, upload-time = "2026-05-07T08:03:47.816Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/37/d3/d7dff0d58a9e9244b48044bfb6a898bfcc8ecc42e0031d1bebc695344725/google_auth_oauthlib-1.4.0-py3-none-any.whl", hash = "sha256:251314f213a9ee46a5ae73988e84fd7cca8bb68e7ecf4bfd45940f9e7f51d070", size = 19261, upload-time = "2026-05-07T08:02:13.798Z" }, -] - [[package]] name = "google-cloud-aiplatform" version = "1.154.0" @@ -1133,23 +1069,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1b/bd/a89eaebd2f9db5f92ddcc8e4f23c266be1dbd11058bb83451d8dd029f34c/google_cloud_storage-3.12.0-py3-none-any.whl", hash = "sha256:3880773754ddf7c27567b04e2a4d193950b6b99429f37b9097d873686e95b09c", size = 340605, upload-time = "2026-06-12T18:03:12.677Z" }, ] -[[package]] -name = "google-cloud-storage-control" -version = "1.12.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "google-api-core", extra = ["grpc"] }, - { name = "google-auth" }, - { name = "grpc-google-iam-v1" }, - { name = "grpcio" }, - { name = "proto-plus" }, - { name = "protobuf" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/88/ae/707995271f77e3c1da715c740160f98f87a79a5c601b28d17c4d2500c037/google_cloud_storage_control-1.12.0.tar.gz", hash = "sha256:49090d03532c0c84c6246a5fd490c31fd4bbffb4c7771c0a6744ec23a194a5f6", size = 137036, upload-time = "2026-06-03T16:14:00.921Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/34/ba/f187f30fb2c8bb5dbb0de0fcb0dec2f8fab0b782ab42facc83ea02628f1b/google_cloud_storage_control-1.12.0-py3-none-any.whl", hash = "sha256:20f7f1252fa5635d47e12f746aa69d719e31eb9405cbbd14cdfba317db27d421", size = 102205, upload-time = "2026-06-03T16:12:43.266Z" }, -] - [[package]] name = "google-crc32c" version = "1.8.0" @@ -1239,57 +1158,57 @@ wheels = [ [[package]] name = "grpcio" -version = "1.82.0" +version = "1.83.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b6/19/e29d3979b420b92d516ee97f0bff4fb88cbb3d724791318f50beb55db549/grpcio-1.82.0.tar.gz", hash = "sha256:bfe3247e0bb598585ce26565730970d21bccf3c9dc547dc6703a1a3c7888e8e1", size = 13184479, upload-time = "2026-07-06T04:22:54.421Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2f/e6/9ce68bc431224177b6f506cfb4896a742119c51255143069970955b6510b/grpcio-1.82.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:03cce11532da292215cd8e5f444de91982e39dfb41a916e0e0f91a5c128588ea", size = 6144680, upload-time = "2026-07-06T04:21:29.808Z" }, - { url = "https://files.pythonhosted.org/packages/59/93/00ecf28e1be7a70b4b4d148d3a551b6c9a37024a669c3650a2c374c64026/grpcio-1.82.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:01336fe31d1ea5d5154b0d803eb3584e69fb96e0e657acc87bd4d48d6abc9587", size = 11952213, upload-time = "2026-07-06T04:21:32.347Z" }, - { url = "https://files.pythonhosted.org/packages/38/05/9be6bf4cddb5b5f39c0748cdf5639525cd4116a157c8f3802c9217e54882/grpcio-1.82.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fe56f1bea709ddb2531fbd510a2dd6040e993985507c3b2bf3404507aee989b1", size = 6710770, upload-time = "2026-07-06T04:21:35.382Z" }, - { url = "https://files.pythonhosted.org/packages/00/97/62c0969e993673ebef16d526db2504f38bc4c73cf2593c28746791ab7956/grpcio-1.82.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:77a5d8fe331376c7aaa4e06e903a43bd144de4f98c6e414e13cbc5d64e5aa61e", size = 7450671, upload-time = "2026-07-06T04:21:37.923Z" }, - { url = "https://files.pythonhosted.org/packages/67/1f/d83d9fbaaef2b6ebbf70a6e467000f13f92b4cf0b84c561c162b43128439/grpcio-1.82.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9ce80ea66c803398fde9c5b1fd925920c9742da7f10f8b85acac0adac3e4d7e0", size = 6886855, upload-time = "2026-07-06T04:21:40.527Z" }, - { url = "https://files.pythonhosted.org/packages/c5/36/5ed2e28b8c5f00599c7d5e94d5f82c32cf73a6fa42d13c2900cb37c861ec/grpcio-1.82.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4665dc8c9ec30acff455e2b15c47b825f18647c555483aa1ccf670adead8adcc", size = 7501322, upload-time = "2026-07-06T04:21:43.78Z" }, - { url = "https://files.pythonhosted.org/packages/89/4e/69113709e14e24289940d6aef067b797c73c8c96f580683af7d5f09b22b0/grpcio-1.82.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4405d19ecfc7a8caa9043ff5a651e1871b54fc620f917de5dede7e6fb78e9e14", size = 8536896, upload-time = "2026-07-06T04:21:46.387Z" }, - { url = "https://files.pythonhosted.org/packages/30/f6/08bd6eb89a558f0f59dacaef747afda7c21e2c2ecf138ae2ec533dea8aeb/grpcio-1.82.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:30e580c19178609799660f52b1e1eb09248969be20dc44caaa4dcaa3e8450dd9", size = 7913890, upload-time = "2026-07-06T04:21:49.733Z" }, - { url = "https://files.pythonhosted.org/packages/d5/a1/04b467e2ffc93f8e50b660bc6f47c3c8e9ba7603104a741e3d5663a261d4/grpcio-1.82.0-cp312-cp312-win32.whl", hash = "sha256:dbbd83dbaa856b387a4eba74ec3add7f594f090d3b4d390d829dd5834816aae4", size = 4240969, upload-time = "2026-07-06T04:21:52.281Z" }, - { url = "https://files.pythonhosted.org/packages/25/3c/15d3601384d5937e8d9fbe6d8e7b8171008d649744dd7b6c864932937ec9/grpcio-1.82.0-cp312-cp312-win_amd64.whl", hash = "sha256:73176d270699d76a9dc3753f25e01c19fb98f830f826a506e917d83629f1b39b", size = 5001575, upload-time = "2026-07-06T04:21:54.839Z" }, - { url = "https://files.pythonhosted.org/packages/f0/c8/7aace81e7cd12c6ffccf0d47ac389a3f19e5280e9ab04d301d02855ecd25/grpcio-1.82.0-cp313-cp313-linux_armv7l.whl", hash = "sha256:841bcaa24882921d41cd7a82118f9a766edcf8807c2f486e7deeb0ff5ea36181", size = 6146066, upload-time = "2026-07-06T04:21:57.527Z" }, - { url = "https://files.pythonhosted.org/packages/10/17/b448f2265e4927188425c0b09fb943c39b50f7f53fbead644b44ccdf834b/grpcio-1.82.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:d4bb7219c39c735e6573ae3c33aeafb2a4e1f1cc93a762f7899c283defaed365", size = 11948617, upload-time = "2026-07-06T04:22:00.066Z" }, - { url = "https://files.pythonhosted.org/packages/8e/a3/fba54e50fffc9f74b1ffd5eb4e47310e6650557ef136a165f1fe7df03a65/grpcio-1.82.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:254230d2cb773a87380858d0fc74653b4761bf2013c83fe1c5d77ed8f241fc9d", size = 6714591, upload-time = "2026-07-06T04:22:04.04Z" }, - { url = "https://files.pythonhosted.org/packages/29/6f/f0ec3f1a61a746dc7037a737ba4cbe60959645311ac8542bb1ad4e72ae8b/grpcio-1.82.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:16d4f215344057eb21038319b05a0c531306230c17b075fa8461a944581006dd", size = 7454986, upload-time = "2026-07-06T04:22:06.776Z" }, - { url = "https://files.pythonhosted.org/packages/b6/11/9d6ce94465d6a7f92c413430b3e77f0879b39427a217ffb3d23585df4bb3/grpcio-1.82.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a7c283b459151a2e84df32b28c31562ab3e7f624bccffce176c5608565b0212f", size = 6888623, upload-time = "2026-07-06T04:22:09.469Z" }, - { url = "https://files.pythonhosted.org/packages/51/1a/b887adb7abb69081c4f24bc66c3612975e87caf5f7c6aa3916bb82d42e5f/grpcio-1.82.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:96a3be4d2a482ff004c036f6391f33553b36a7627138d1ca3ecc1e70d55d8af5", size = 7505067, upload-time = "2026-07-06T04:22:11.728Z" }, - { url = "https://files.pythonhosted.org/packages/60/94/90f11c0f227ac653cff769e05c7f279da945e134c9351a1c37d475714686/grpcio-1.82.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:bc481e6c7e6c8721797f2ab764092c843fd5cea06d4c7e9f4c3e3b7ecf331eb0", size = 8535382, upload-time = "2026-07-06T04:22:14.453Z" }, - { url = "https://files.pythonhosted.org/packages/d2/3d/de8569386685213135dfbaab58e221add11858485524cf7d5987efc6d70b/grpcio-1.82.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:680f633bc788652222cb90a568ec374edcd91e1fe7715a7c248ece8e1032c2bd", size = 7910708, upload-time = "2026-07-06T04:22:17.994Z" }, - { url = "https://files.pythonhosted.org/packages/42/ab/3950b32369763818ae89ef76d52e84bd034d3a0ab46bf315669c913b32e2/grpcio-1.82.0-cp313-cp313-win32.whl", hash = "sha256:640b4715b9c206e5176e46601aa1422c47c779f642a869dfd27062e8fde13dfb", size = 4240351, upload-time = "2026-07-06T04:22:20.626Z" }, - { url = "https://files.pythonhosted.org/packages/a7/5c/c7b9a48efe973883f5707a311f87772a4c307a22ec80d6f3c851846a0a02/grpcio-1.82.0-cp313-cp313-win_amd64.whl", hash = "sha256:8523e81bd23f83e607aed28fbd8244b2be4aeb34e084fcef1bd1867709085c2a", size = 5000985, upload-time = "2026-07-06T04:22:23.365Z" }, - { url = "https://files.pythonhosted.org/packages/bd/6e/abc5efe11b73eefa0183531032eef2a53f33aeced6484b0d9da559c12ed6/grpcio-1.82.0-cp314-cp314-linux_armv7l.whl", hash = "sha256:136969867c14fcc743fa39b12d7da133cae7da4f8e0e7767b6b8b7e75a8c24fd", size = 6146898, upload-time = "2026-07-06T04:22:26.082Z" }, - { url = "https://files.pythonhosted.org/packages/e0/e4/d01ee88ceac221436dce1584ae1f046bd44d0dffc5a92e95aa34857c4516/grpcio-1.82.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:a7e246cd216662f513ae79ea3afadde6afe8a659b48a4170bae9a896e69ac254", size = 11954897, upload-time = "2026-07-06T04:22:29.114Z" }, - { url = "https://files.pythonhosted.org/packages/99/b9/2a35540ec08afc08859c35afd1a8a51e2585f39b12cf4eee68dcb1fa973d/grpcio-1.82.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:85c2a4e9e8a30d73d41e547a5101e57a1ed2093ec4b2daf6847c344adba00523", size = 6723102, upload-time = "2026-07-06T04:22:32.146Z" }, - { url = "https://files.pythonhosted.org/packages/ec/ba/af71af59943b5d08879617a1d97495d45c53e43f3c6ef16f820c64e107f9/grpcio-1.82.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:a3a7802328e0d20e6ac458ff1974a8096cfaa3ca84dea0903235f641f1d703c1", size = 7454545, upload-time = "2026-07-06T04:22:34.776Z" }, - { url = "https://files.pythonhosted.org/packages/b3/7e/258aae12b68910140725873657469ff166f8968bdb0674e12d6452b1812c/grpcio-1.82.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9331e6b91b26cffd63fb70a545f3c2cac8dc1b434b4436f43915c23e1e22f265", size = 6889585, upload-time = "2026-07-06T04:22:37.411Z" }, - { url = "https://files.pythonhosted.org/packages/5e/04/cfada8930306b9101cfb405b33ecb3b72abce51d725f900f167ffbc3146b/grpcio-1.82.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:670ab93a0059e707de53b9e715304b9f566b6defee80e5490cdd15820ce032f6", size = 7514162, upload-time = "2026-07-06T04:22:40.076Z" }, - { url = "https://files.pythonhosted.org/packages/2e/57/955f95989338857a6d73ee838b4d9e8198d50972a70ca83ec22322396f4c/grpcio-1.82.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:fbfe2e321bb30b84799677d6e8469896e487b494e9e4fbc9b8cc4425fe054d44", size = 8536163, upload-time = "2026-07-06T04:22:42.801Z" }, - { url = "https://files.pythonhosted.org/packages/06/62/826da0b0f01c7ba3f7ea0185968f551290deb28968d98b1edf604c895db8/grpcio-1.82.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2472b72a55da8570e089f1ca22f34a350d86c109be2755f6ee1cca5ebbdd9b54", size = 7912573, upload-time = "2026-07-06T04:22:46.187Z" }, - { url = "https://files.pythonhosted.org/packages/4d/b4/d8824a66ac1ac7fbcdf63806f1a8725d95426654e859afa4606070c78740/grpcio-1.82.0-cp314-cp314-win32.whl", hash = "sha256:c9feb986151c2726789599b3672c1c9faa1d0824da921dc3f33a8cb1f93e2861", size = 4321824, upload-time = "2026-07-06T04:22:48.866Z" }, - { url = "https://files.pythonhosted.org/packages/da/08/02e581cc0bbb4c7b842f85e66a945b46d5bcc5927575c1086edfbc3360e7/grpcio-1.82.0-cp314-cp314-win_amd64.whl", hash = "sha256:b1eeb0480d86965263f8c8ae7ee5360c284d22176a196aab02fce7fba8d89dd8", size = 5141112, upload-time = "2026-07-06T04:22:51.443Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/e2/b1/46539f5050d7c316a13396d185451f95084a74ddc68b12d818595bef0377/grpcio-1.83.1.tar.gz", hash = "sha256:9cee6fcbf2eb57c4b49451787bfa87be8efc1ca02a0b327dd4b54d44502e362b", size = 13445033, upload-time = "2026-08-28T07:09:11.464Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/85/9e/a3ba13e08bbee5bf6e57597dfe4823961fd7e94c0b8afe3a4bb7dca639f3/grpcio-1.83.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:5acd14c6ddf047de62cbf8745b11103ea91abbf57d1b8edd5395ccd9fcd13abb", size = 6303170, upload-time = "2026-08-28T07:08:08.188Z" }, + { url = "https://files.pythonhosted.org/packages/1c/ae/65ce56a2527faa17d02cba4c2231c74047ad898be339486ba87f093bfb66/grpcio-1.83.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:16138031a47b771860a16a975b53087f4fd5bbdbb2c03a188c5d90ad65d2bdae", size = 12165806, upload-time = "2026-08-28T07:08:10.309Z" }, + { url = "https://files.pythonhosted.org/packages/4e/91/40432480088a2243d360864de072ed5b78c4ebbaabd29c28918f1e1b1454/grpcio-1.83.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5ccc26715fd4defca5e129e280dd883b1737b65045ec50ffe22ce42104089519", size = 6872490, upload-time = "2026-08-28T07:08:12.355Z" }, + { url = "https://files.pythonhosted.org/packages/c8/62/3da2300c8c79fd20a78a8a4bb6251e5068d9af33bc8fd389b98fec35e8a3/grpcio-1.83.1-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:b74f2a1d9ab1dfa3e263ef33d581613679b78d0884babf11671af26e45570ead", size = 7618367, upload-time = "2026-08-28T07:08:14.025Z" }, + { url = "https://files.pythonhosted.org/packages/bc/19/9fc702e31a631262d7a752fa699f6022821e707fefc8bff49b1550a57729/grpcio-1.83.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:72578aa07a4008f17521ef52debcc3acfd1e2c5426243bc3ffb56a38bfe610b7", size = 7040936, upload-time = "2026-08-28T07:08:15.963Z" }, + { url = "https://files.pythonhosted.org/packages/ec/56/95933cc44cba2429765fa065c951dd529e5771b119d9d2481b4646f1d6a5/grpcio-1.83.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c12e1fc59c6dc26d10d9144453ddc6cbfe4cd4c31e874ed2d0132f88e685eb8b", size = 7573096, upload-time = "2026-08-28T07:08:17.729Z" }, + { url = "https://files.pythonhosted.org/packages/ac/80/af63359da06b016de48cb111f144703a10043850dafa43ae0a038907b9e8/grpcio-1.83.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4910b62f7d12197160bfb7de06d876d64dd12d43483e8292f98f49ca09b628d9", size = 8609442, upload-time = "2026-08-28T07:08:19.777Z" }, + { url = "https://files.pythonhosted.org/packages/3d/fa/f0586c56bdfb8a7a2adda01e0ac2413447cde3141ab09411a5d5afdcffd3/grpcio-1.83.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:9e703effe3ae779925c82ac24fdb82cf4105e1096810151ed9501c5f34546b9c", size = 7984321, upload-time = "2026-08-28T07:08:22.114Z" }, + { url = "https://files.pythonhosted.org/packages/25/8a/14ec05669f9eb295801e26c2ea8c561a1b786b0e3557c2c22131165ab010/grpcio-1.83.1-cp312-cp312-win32.whl", hash = "sha256:a2aea8bd6e0a34f12cbaddb7bb70bec836818789fa5c7ab7572c6b745396a2d4", size = 4395604, upload-time = "2026-08-28T07:08:24.08Z" }, + { url = "https://files.pythonhosted.org/packages/e9/37/8c2f7cc16089e36a3fbacaacc7a3d043912aa0d2dfae5556f6450414ea6e/grpcio-1.83.1-cp312-cp312-win_amd64.whl", hash = "sha256:583bf2e8255040a4a312f9572dfe62a05271437b149550e1a536d5c47d2d1e8a", size = 5161512, upload-time = "2026-08-28T07:08:25.81Z" }, + { url = "https://files.pythonhosted.org/packages/7c/fd/d1fc58933bf88c9209f89dc570c810f1aa57cb04b3459cf2b26f61e32112/grpcio-1.83.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:8d228e253b77865efcbdd7b5894ca882c9e0ea98c02b7d20582e61ded8dfd4b5", size = 6305628, upload-time = "2026-08-28T07:08:27.872Z" }, + { url = "https://files.pythonhosted.org/packages/c4/49/0b40bae059c619505c9b751cee6caa208e4904e290aaefa1728c4c2c67a5/grpcio-1.83.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:0468b627f2987c9a77f7580030207cbd85457ffe52998beff4f0b5c38c58a72c", size = 12156839, upload-time = "2026-08-28T07:08:30.191Z" }, + { url = "https://files.pythonhosted.org/packages/61/4b/e8c0d635da0ee5ddd9950c8d540f5dcdd0ef1854a382cc55496a487a8d31/grpcio-1.83.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a6a282e81530cead60bbd752cc04950a57f224379e9821495d6a35bd5ce9b1f4", size = 6877036, upload-time = "2026-08-28T07:08:32.285Z" }, + { url = "https://files.pythonhosted.org/packages/c5/d4/760a33f339a7dd3d5f4b3e0e9bec5472d95592a80f887b2e9dab4e41cfbc/grpcio-1.83.1-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:947d945f52e8ecf3cafd2bb7113502a16ccfda3e12c854443094de32d83ad432", size = 7624404, upload-time = "2026-08-28T07:08:34.194Z" }, + { url = "https://files.pythonhosted.org/packages/54/ec/bd798654b06fb42a92b57d1dc1b530084fa89ed442806fcd0a833a36f9b3/grpcio-1.83.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:55656318d5dd387077396dffb929171ca3966e24bfead9a6c5dba9f889062cb4", size = 7042942, upload-time = "2026-08-28T07:08:36.208Z" }, + { url = "https://files.pythonhosted.org/packages/08/b0/c00f86614566dd0961825cf0f43d4f96a74371d9d95f952bcbc4b86d9a27/grpcio-1.83.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:9daf5acf4fc9d5f5627229969c2580a91e511779d76e4ccdeb9f4770f05d8bc2", size = 7576937, upload-time = "2026-08-28T07:08:38.041Z" }, + { url = "https://files.pythonhosted.org/packages/b1/38/85eff43a5c89dc666a252b5c9f8e9ab03f89e11c95b6263d2933f08fdbe7/grpcio-1.83.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:7b94174cbca93316888f805efbeb08f1c020f7b7493d2d50cc4f6b64ebb7e8bd", size = 8608391, upload-time = "2026-08-28T07:08:40.092Z" }, + { url = "https://files.pythonhosted.org/packages/35/4e/82835483e2f812494be865e7965c0d626cb9e71ab0d83a420d75aea4ad67/grpcio-1.83.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:65c5a7210911ffe0f67b1cdc5308f9854b6d1f1b345e3e49ab7cac1ba50fa346", size = 7980060, upload-time = "2026-08-28T07:08:42.434Z" }, + { url = "https://files.pythonhosted.org/packages/b2/b7/68a98bef733fef704fbcfb3957c8dba67e3e38ca7a7fea851195bc97c648/grpcio-1.83.1-cp313-cp313-win32.whl", hash = "sha256:179368d9361854616ce6f397d4716e07480129652752fcbcfc5a7260455ad6f2", size = 4395226, upload-time = "2026-08-28T07:08:44.463Z" }, + { url = "https://files.pythonhosted.org/packages/85/a0/df4de3b51d37ac8fb0320bb9668381ce2bd3b7aa990880bfc56a8a26f665/grpcio-1.83.1-cp313-cp313-win_amd64.whl", hash = "sha256:2e57af456385491a76e13c4aada8c8f43a8e47051e06ea97a9dbe2a49654e6db", size = 5160273, upload-time = "2026-08-28T07:08:46.216Z" }, + { url = "https://files.pythonhosted.org/packages/42/9c/484d981d8b90c4e6abf3030bd2ed747e84d1eb192b3ec9cbb41e0b73e4bf/grpcio-1.83.1-cp314-cp314-linux_armv7l.whl", hash = "sha256:8b3c87ca908296bf125f841d3e1a2225a2b39aaa8ed7a57e7ccde465ee519bab", size = 6306089, upload-time = "2026-08-28T07:08:48.379Z" }, + { url = "https://files.pythonhosted.org/packages/84/01/0afec1c92e4f292f74a44ecf75eabbf40903125b8c4df103c9868d6338da/grpcio-1.83.1-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:c0f3f20c90e72a171917ae65706500b096a1c3eb5f162c3ce702a2e25635f132", size = 12170381, upload-time = "2026-08-28T07:08:50.653Z" }, + { url = "https://files.pythonhosted.org/packages/c7/5a/e9a2383804433a0a61d6d93777ad321c7f36ac1cfdaa4c6d1a7c9ac846b7/grpcio-1.83.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:81bbf35a46bf8cad2dfbb2eccc19c711befb58b288acb534bbcd0d74283202a6", size = 6883286, upload-time = "2026-08-28T07:08:53.654Z" }, + { url = "https://files.pythonhosted.org/packages/63/e7/f8ca8f76994e14c70b9a0052e82f10de497a23db450c36379c9716ebfc4d/grpcio-1.83.1-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:215cec07d11176507387bda4bf2751816e880f9bff8dc1ca524bfbb8ed8f2fad", size = 7624293, upload-time = "2026-08-28T07:08:55.709Z" }, + { url = "https://files.pythonhosted.org/packages/bb/7a/4b672814b0cd0fe63bdd735379d88b165759f3144ab023ad8ec5fc4d53ac/grpcio-1.83.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:abce7d43ec29cd39230fa8339de1a07643b55adc412a454850fbd875349950ff", size = 7044346, upload-time = "2026-08-28T07:08:57.802Z" }, + { url = "https://files.pythonhosted.org/packages/50/b8/d89fe60e4239ad51be333dd9cc703741d449a35064e51f8a0b5bfa755432/grpcio-1.83.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:e256f95a40e3b0183a98556fb7164d24b97eeb353123ccabfcba94712b35ee2a", size = 7584187, upload-time = "2026-08-28T07:08:59.867Z" }, + { url = "https://files.pythonhosted.org/packages/dc/b2/b290d7402633d9166e4dd47e6f5f74a24ce10a8340b84455896ebc349f85/grpcio-1.83.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:2110059146fb0ea216e1ffddb29377b5cc2fd412a5b0a92e102616bd5edf18c2", size = 8608730, upload-time = "2026-08-28T07:09:02.592Z" }, + { url = "https://files.pythonhosted.org/packages/f5/44/fa89e44d1b5cf5b9fa71b2fd7abf506f182fd43917231a92fbf1ea326b02/grpcio-1.83.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:20d944d967843f8183f9f23d5916388362e5f8eeeae855bbe4354d906dc9f31b", size = 7983283, upload-time = "2026-08-28T07:09:05.087Z" }, + { url = "https://files.pythonhosted.org/packages/ce/b8/9db73ed1f35ffa76124ac574bf296d06a359798dfd6b50d382f2b8a060a1/grpcio-1.83.1-cp314-cp314-win32.whl", hash = "sha256:623c87c6d4a1cb30d82c4e896f95477050f2e01b4a1f8cf91ff2b1abdf89c457", size = 4474327, upload-time = "2026-08-28T07:09:07.179Z" }, + { url = "https://files.pythonhosted.org/packages/65/22/fc9a622d885a7a37ff972a12faaef443d74e47407181da70d0ab62ab41f0/grpcio-1.83.1-cp314-cp314-win_amd64.whl", hash = "sha256:47e6934ad38779271e2e7cc5f78a63a407cf3d98114c65c1fdbcd3f5a716f29b", size = 5302032, upload-time = "2026-08-28T07:09:09.285Z" }, ] [[package]] name = "grpcio-status" -version = "1.82.0" +version = "1.83.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "googleapis-common-protos" }, { name = "grpcio" }, { name = "protobuf" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/07/c3/42148fdb28747b6c12df1351dece3fd574262fa14b180eca9821428626df/grpcio_status-1.82.0.tar.gz", hash = "sha256:ea0cbc12c52a902cb33ba56fb3233dcfad53ffead08f4a659b1752d473a6a0ac", size = 13902, upload-time = "2026-07-06T04:26:47.864Z" } +sdist = { url = "https://files.pythonhosted.org/packages/db/03/dee41fe15a9657c60397c1f215622a1f146e174367bbbb67911b62ee0629/grpcio_status-1.83.1.tar.gz", hash = "sha256:c08c8d553d6ab96effae1d923839de22926f5a316a942f48a48040e047c517af", size = 13957, upload-time = "2026-08-28T07:12:25.893Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ad/f3/fe7e1f57f170ed6cae20eecc0dea62e9de862e68adede80a6cd54a9b4dfa/grpcio_status-1.82.0-py3-none-any.whl", hash = "sha256:d41e45164ae2c265b268aebd9d86defa292a1ba91d86a548bad1f059c4645676", size = 14634, upload-time = "2026-07-06T04:26:38.672Z" }, + { url = "https://files.pythonhosted.org/packages/e7/34/b03ec688e5c8c6ce283ec8b214d9c171e97059eb55157465f6fd6db1562c/grpcio_status-1.83.1-py3-none-any.whl", hash = "sha256:2cd328ee62ef2b3eb957dd3b75db7dadcdbb76488fdc9ab3aba1ebfbbdc324a4", size = 14636, upload-time = "2026-08-28T07:12:15.982Z" }, ] [[package]] @@ -1996,6 +1915,46 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5e/75/bd9b7bb966668920f06b200e84454c8f3566b102183bc55c5473d96cb2b9/msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca", size = 20583, upload-time = "2025-03-14T23:51:03.016Z" }, ] +[[package]] +name = "msgspec" +version = "0.21.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/60/f79b9b013a16fa3a58350c9295ddc6789f2e335f36ea61ed10a21b215364/msgspec-0.21.1.tar.gz", hash = "sha256:2313508e394b0d208f8f56892ca9b2799e2561329de9763b19619595a6c0f72c", size = 319193, upload-time = "2026-04-12T21:44:50.394Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6e/cf/317224852c00248c620a9bcf4b26e2e4ab8afd752f18d2a6ef73ebd423b6/msgspec-0.21.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d4248cf0b6129b7d230eacd493c17cc2d4f3989f3bb7f633a928a85b7dcfa251", size = 196188, upload-time = "2026-04-12T21:44:07.181Z" }, + { url = "https://files.pythonhosted.org/packages/6d/81/074612945c0666078f7366f40000013de9f6ba687491d450df699bceebc9/msgspec-0.21.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5102c7e9b3acff82178449b85006d96310e690291bb1ea0142f1b24bcb8aabcb", size = 188473, upload-time = "2026-04-12T21:44:08.736Z" }, + { url = "https://files.pythonhosted.org/packages/8a/37/655101799590bcc5fddb2bd3fe0e6194e816c2d1da7c361725f5eb89a910/msgspec-0.21.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:846758412e9518252b2ac9bffd6f0e54d9ff614f5f9488df7749f81ff5c80920", size = 218871, upload-time = "2026-04-12T21:44:09.917Z" }, + { url = "https://files.pythonhosted.org/packages/b5/d1/d4cd9fe89c7d400d7a18f86ccc94daa3f0927f53558846fcb60791dce5d6/msgspec-0.21.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:21995e74b5c598c2e004110ad66ec7f1b8c20bf2bcf3b2de8fd9a3094422d3ff", size = 225025, upload-time = "2026-04-12T21:44:11.191Z" }, + { url = "https://files.pythonhosted.org/packages/24/bf/e20549e602b9edccadeeff98760345a416f9cce846a657e8b18e3396b212/msgspec-0.21.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6129f0cca52992e898fd5344187f7c8127b63d810b2fd73e36fca73b4c6475ee", size = 222672, upload-time = "2026-04-12T21:44:12.481Z" }, + { url = "https://files.pythonhosted.org/packages/b4/68/04d7a8f0f786545cf9b8c280c57aa6befb5977af6e884b8b54191cbe44b3/msgspec-0.21.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ef3ec2296248d1f8b9231acb051b6d471dfde8f21819e86c9adaaa9f42918521", size = 227303, upload-time = "2026-04-12T21:44:13.709Z" }, + { url = "https://files.pythonhosted.org/packages/cc/4d/619866af2840875be408047bf9e70ceafbae6ab50660de7134ed1b25eb86/msgspec-0.21.1-cp312-cp312-win_amd64.whl", hash = "sha256:d4ab834a054c6f0cbeef6df9e7e1b33d5f1bc7b86dea1d2fd7cad003873e783d", size = 190017, upload-time = "2026-04-12T21:44:14.977Z" }, + { url = "https://files.pythonhosted.org/packages/5e/2e/a8f9eca8fd00e097d7a9e99ba8a4685db994494448e3d4f0b7f6e9a3c0f7/msgspec-0.21.1-cp312-cp312-win_arm64.whl", hash = "sha256:628aaa35c74950a8c59da330d7e98917e1c7188f983745782027748ee4ca573e", size = 175345, upload-time = "2026-04-12T21:44:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/7e/74/f11ede02839b19ff459f88e3145df5d711626ca84da4e23520cebf819367/msgspec-0.21.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:764173717a01743f007e9f74520ed281f24672c604514f7d76c1c3a10e8edb66", size = 196176, upload-time = "2026-04-12T21:44:17.613Z" }, + { url = "https://files.pythonhosted.org/packages/bb/40/4476c1bd341418a046c4955aff632ec769315d1e3cb94e6acf86d461f9ed/msgspec-0.21.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:344c7cd0eaed1fb81d7959f99100ef71ec9b536881a376f11b9a6c4803365697", size = 188524, upload-time = "2026-04-12T21:44:18.815Z" }, + { url = "https://files.pythonhosted.org/packages/ca/d9/9e9d7d7e5061b47540d03d640fab9b3965ba7ae49c1b2154861c8f007518/msgspec-0.21.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:48943e278b3854c2f89f955ddc6f9f430d3f0784b16e47d10604ee0463cd21f5", size = 218880, upload-time = "2026-04-12T21:44:20.028Z" }, + { url = "https://files.pythonhosted.org/packages/74/66/2bb344f34abb4b57e60c7c9c761994e0417b9718ec1460bf00c296f2a7ea/msgspec-0.21.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a9aa659ebb0101b1cbc31461212b87e341d961f0ab0772aaf068a99e001ec4aa", size = 225050, upload-time = "2026-04-12T21:44:21.577Z" }, + { url = "https://files.pythonhosted.org/packages/1a/84/7c1e412f76092277bf760cef12b7979d03314d259ab5b5cafde5d0c1722d/msgspec-0.21.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7b27d1a8ead2b6f5b0c4f2d07b8be1ccfcc041c8a0e704781edebe3ae13c484", size = 222713, upload-time = "2026-04-12T21:44:22.83Z" }, + { url = "https://files.pythonhosted.org/packages/4e/27/0bba04b2b4ef05f3d068429410bc71d2cea925f1596a8f41152cccd5edb8/msgspec-0.21.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:38fe93e86b61328fe544cb7fd871fad5a27c8734bfda90f65e5dbe288ae50f61", size = 227259, upload-time = "2026-04-12T21:44:24.11Z" }, + { url = "https://files.pythonhosted.org/packages/b0/2d/09574b0eea02fed2c2c1383dbaae2c7f79dc16dcd6487a886000afb5d7c4/msgspec-0.21.1-cp313-cp313-win_amd64.whl", hash = "sha256:8bc666331c35fcce05a7cd2d6221adbe0f6058f8e750711413d22793c080ac6a", size = 189857, upload-time = "2026-04-12T21:44:25.359Z" }, + { url = "https://files.pythonhosted.org/packages/46/34/105b1576ad182879914f0c821f17ee1d13abb165cb060448f96fe2aff078/msgspec-0.21.1-cp313-cp313-win_arm64.whl", hash = "sha256:42bb1241e0750c1a4346f2aa84db26c5ffd99a4eb3a954927d9f149ff2f42898", size = 175403, upload-time = "2026-04-12T21:44:26.608Z" }, + { url = "https://files.pythonhosted.org/packages/5a/ad/86954e987d1d6a5c579e2c2e7832b65e0fff194179fdac4f581536086024/msgspec-0.21.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:fab48eb45fdbfbdb2c0edfec00ffc53b6b6085beefc6b50b61e01659f9f8757f", size = 196261, upload-time = "2026-04-12T21:44:27.807Z" }, + { url = "https://files.pythonhosted.org/packages/d1/a1/c5e46c3e42b866199365e35d11dddfd1fbd8bba4fdb3c52f965b1607ce94/msgspec-0.21.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:3cb779ea0c35bc807ff941d415875c1f69ca0be91a2e907ab99a171811d86a9a", size = 188729, upload-time = "2026-04-12T21:44:28.99Z" }, + { url = "https://files.pythonhosted.org/packages/85/7d/1e29a319d678d6cb962ae5bdf32a6858ebdf38f73bc654c0e9c742a0c2c8/msgspec-0.21.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:68604db36b3b4dd9bf160e436e12798a4738848144cea1aca1cb984011eb160f", size = 219866, upload-time = "2026-04-12T21:44:31.104Z" }, + { url = "https://files.pythonhosted.org/packages/25/1f/cca084ca2572810fff12ea9dbdcbe39eac048f40daf4a9077b49fcbe8cee/msgspec-0.21.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3d6b9dc50948eaf65df54d2fd0ff66e6d8c32f116037209ee861810eb9b676cb", size = 224993, upload-time = "2026-04-12T21:44:32.649Z" }, + { url = "https://files.pythonhosted.org/packages/71/94/d2120fc9d419a89a3a7c13e5b7078798c4b392a96a02a6e2b3ce43a8766c/msgspec-0.21.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:52c5e21930942302394429c5a582ce7e6b62c7f983b3760834c2ce107e0dd6df", size = 223535, upload-time = "2026-04-12T21:44:33.839Z" }, + { url = "https://files.pythonhosted.org/packages/75/17/42418b66a3ad972a89bab73dd78b79cc6282bb488a25e73c853cee7443b9/msgspec-0.21.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:abbb39d65681fa24ed394e01af3d59d869068324f900c61d06062b7fb9980f2f", size = 227222, upload-time = "2026-04-12T21:44:35.093Z" }, + { url = "https://files.pythonhosted.org/packages/c4/33/265c894268cca88ff67b144ca2b4c522fc8b9a6f1966a3640c70516e78e1/msgspec-0.21.1-cp314-cp314-win_amd64.whl", hash = "sha256:5666b1b560b97b6ec2eb3fca8a502298ebac56e13bbca1f88523538ce83d01ea", size = 193810, upload-time = "2026-04-12T21:44:36.612Z" }, + { url = "https://files.pythonhosted.org/packages/3b/8f/a6d35f25bf1fc63c492fdd88fdce01ba0875ead48c2b91f90f33653b4131/msgspec-0.21.1-cp314-cp314-win_arm64.whl", hash = "sha256:d8b8578e4c83b14ceea4cef0d0b747e31d9330fe4b03b2b2ad4063866a178f93", size = 179125, upload-time = "2026-04-12T21:44:38.198Z" }, + { url = "https://files.pythonhosted.org/packages/c6/39/74839641e64b99d87da55af0fc472854d42b46e2183b9e2a67fe1bb2a512/msgspec-0.21.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:15f523d51c00ebad412213bfe9f06f0a50ec2b93e0c19e824a2d267cabb48ea2", size = 200171, upload-time = "2026-04-12T21:44:39.414Z" }, + { url = "https://files.pythonhosted.org/packages/70/9b/ce0cca6d2d87fcd4b6ff97600790494e64f26a2c55d61507cd2755c16193/msgspec-0.21.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:4e47390360583ba3d5c6cb44cf0a9f61b0a06a899d3c2c00627cedebb2e2884b", size = 192879, upload-time = "2026-04-12T21:44:40.882Z" }, + { url = "https://files.pythonhosted.org/packages/a7/08/673a7bb05e5702dc787ddd3011195b509f9867927970da59052211929987/msgspec-0.21.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f60800e6299b798142dc40b0644da77ceac5ea0568be58228417eae14135c847", size = 226281, upload-time = "2026-04-12T21:44:42.181Z" }, + { url = "https://files.pythonhosted.org/packages/7d/45/86508cf57283e9070b3c447e3ab25b792a7a0855a3ea4e0c6d111ac34c97/msgspec-0.21.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5f8e9dfcd98419cf7568808470c4317a3fb30bef0e3715b568730a2b272a20d7", size = 229863, upload-time = "2026-04-12T21:44:43.442Z" }, + { url = "https://files.pythonhosted.org/packages/2c/62/e7c9367cd08d590559faacd711edbae36840342843e669440363f33c7d36/msgspec-0.21.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:92d89dfad13bd1ea640dc3e37e724ed380da1030b272bdf5ecafb983c3ad7c75", size = 230445, upload-time = "2026-04-12T21:44:44.806Z" }, + { url = "https://files.pythonhosted.org/packages/42/b4/c0f54632103846b658a10930025f4de41c8724b5e4805a5f3b395586cb7e/msgspec-0.21.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:0d03867786e5d7ba25d666df4b11320c27170f4aeafcb8e3a8b0a50a4fb742ca", size = 231822, upload-time = "2026-04-12T21:44:46.343Z" }, + { url = "https://files.pythonhosted.org/packages/ea/1d/0d85cc79d0ccf5508e9c846cc66552a6a16bf92abd1dbd8362617f7b35cd/msgspec-0.21.1-cp314-cp314t-win_amd64.whl", hash = "sha256:740fbf1c9d59992ca3537d6fbe9ebbf9eaf726a65fbf31448e0ecbc710697a63", size = 206650, upload-time = "2026-04-12T21:44:47.601Z" }, + { url = "https://files.pythonhosted.org/packages/90/91/56c5d560f20e6c20e9e4f55bd0e458f7f162aa689ee350346c04c48eac0b/msgspec-0.21.1-cp314-cp314t-win_arm64.whl", hash = "sha256:0d2cc73df6058d811a126ac3a8ad63a4dfa210c82f9cf5a004802eaf4712de90", size = 183149, upload-time = "2026-04-12T21:44:48.833Z" }, +] + [[package]] name = "msrest" version = "0.7.1" @@ -2452,45 +2411,24 @@ wheels = [ ] [[package]] -name = "pinecone-client" -version = "6.0.0" +name = "pinecone" +version = "10.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "certifi" }, - { name = "pinecone-plugin-interface" }, - { name = "python-dateutil" }, - { name = "typing-extensions" }, - { name = "urllib3", marker = "python_full_version < '4'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/6c/ab/3ab3b81e8ad82fbfcaa4f446c7f962b18968d61543c8c9e2c38bd777c056/pinecone_client-6.0.0.tar.gz", hash = "sha256:f224fc999205e4858c4737c40922bdf42d178b361c8859bc486ec00d45b359a9", size = 7004, upload-time = "2025-02-21T19:47:38.899Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5a/e4/7780cd631dc6dad0172a245e958b41b28a70779594c0790fa08b952aa97f/pinecone_client-6.0.0-py3-none-any.whl", hash = "sha256:d81a9e73cae441e4ab6dfc9c1d8b51c9895dae2488cda64f3e21b9dfc10c8d94", size = 6654, upload-time = "2025-02-21T19:47:37.23Z" }, -] - -[[package]] -name = "pinecone-datasets" -version = "0.6.1" -source = { git = "https://github.com/mrT23/pinecone-datasets.git?rev=main#9018d935d27e9bcc2b8de4c151228c60c6cc10da" } -dependencies = [ - { name = "fsspec", version = "2025.12.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, - { name = "fsspec", version = "2026.6.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, - { name = "gcsfs", version = "2025.12.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, - { name = "gcsfs", version = "2026.7.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, - { name = "pandas" }, - { name = "pinecone-client" }, - { name = "pyarrow" }, - { name = "pydantic" }, - { name = "s3fs" }, - { name = "tqdm" }, + { name = "anyio" }, + { name = "httpx", extra = ["http2"] }, + { name = "msgspec" }, + { name = "orjson" }, ] - -[[package]] -name = "pinecone-plugin-interface" -version = "0.0.7" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f4/fb/e8a4063264953ead9e2b24d9b390152c60f042c951c47f4592e9996e57ff/pinecone_plugin_interface-0.0.7.tar.gz", hash = "sha256:b8e6675e41847333aa13923cc44daa3f85676d7157324682dc1640588a982846", size = 3370, upload-time = "2024-06-05T01:57:52.093Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a5/97/f04b682e5d65d4cd09b76afd691f8834cf2c28d82bc942ad637a2c15c74d/pinecone-10.0.0.tar.gz", hash = "sha256:6066bbe9a7ae1d667cde08d262deb6fbea6feb35deb9177dd47141b55bbd9833", size = 525129, upload-time = "2026-09-03T14:31:22.729Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3b/1d/a21fdfcd6d022cb64cef5c2a29ee6691c6c103c4566b41646b080b7536a5/pinecone_plugin_interface-0.0.7-py3-none-any.whl", hash = "sha256:875857ad9c9fc8bbc074dbe780d187a2afd21f5bfe0f3b08601924a61ef1bba8", size = 6249, upload-time = "2024-06-05T01:57:50.583Z" }, + { url = "https://files.pythonhosted.org/packages/0f/c7/9d61f9a4adaaa6e315a0425c2127b449bda64633a9c50bc533ff71c0c894/pinecone-10.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:995c06e905940b10bb2aefe653225340b5a3f56f3efb373fe07f4e57b5043705", size = 3098155, upload-time = "2026-09-03T14:31:11.281Z" }, + { url = "https://files.pythonhosted.org/packages/ec/78/d73c0aa1030931a1b28151c4ff727b793e443a4531521bc89a1d4c3124a0/pinecone-10.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:0994270c514b16c72ec94dd6c29ff2708b81d30ff8467e19de192a28a7c86b7e", size = 3070497, upload-time = "2026-09-03T14:31:12.998Z" }, + { url = "https://files.pythonhosted.org/packages/8a/3e/3b9dba42741ed8faf011d7bce8e1fd10f45dd0503e2888eceafa6be33424/pinecone-10.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d482ed27a805cbd4aca2660da212008dd6e41d87d255279f405ff13b725970e2", size = 3353834, upload-time = "2026-09-03T14:31:14.431Z" }, + { url = "https://files.pythonhosted.org/packages/22/35/50e23cac1d73c921b8cc2eb02c9ad2691f62e8b906db49237bcf85f90676/pinecone-10.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:94d4c64779f3213a5cc538d3bd10a873da192cb9b0039db56690e556ba00b55c", size = 3291380, upload-time = "2026-09-03T14:31:16.052Z" }, + { url = "https://files.pythonhosted.org/packages/04/ef/705e49608eac5803b9b55b6cfae2cac736c0395e709f7668f52c89afc99b/pinecone-10.0.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0e05956a3201b1fbb54a1861277df919318a4941797f2d87fd558ac5ec232151", size = 3528743, upload-time = "2026-09-03T14:31:17.558Z" }, + { url = "https://files.pythonhosted.org/packages/7f/a4/11795c174225b9dea5054618c91b956d0a26b3d7167a5d3531e167712e83/pinecone-10.0.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3ab0c843b4fb04fbac22f1b8455e389063208a50f6a95d7a90627939968198de", size = 3574713, upload-time = "2026-09-03T14:31:19.198Z" }, + { url = "https://files.pythonhosted.org/packages/bd/5d/93d123cfcf458ccc2defa4c0363be5fb7aabfcd375cd2414a3e944e7f55a/pinecone-10.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:2f4e3200ee3562d195802b363487dd7fe6039a8c13630fc25fa3e8726c7a8654", size = 2924160, upload-time = "2026-09-03T14:31:21.107Z" }, ] [[package]] @@ -2554,6 +2492,7 @@ dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp-proto-http" }, { name = "opentelemetry-sdk" }, + { name = "prometheus-client" }, { name = "protobuf" }, { name = "pydantic" }, { name = "pydantic-settings" }, @@ -2596,8 +2535,7 @@ dev = [ similar-issue = [ { name = "lancedb" }, { name = "pandas" }, - { name = "pinecone-client" }, - { name = "pinecone-datasets" }, + { name = "pinecone" }, { name = "qdrant-client" }, ] @@ -2610,7 +2548,7 @@ requires-dist = [ { name = "azure-devops", specifier = "==7.1.0b4" }, { name = "azure-identity", specifier = "==1.25.0" }, { name = "boto3", specifier = "==1.43.78" }, - { name = "certifi", specifier = "==2024.8.30" }, + { name = "certifi", specifier = ">=2026.7.22" }, { name = "dynaconf", specifier = ">=3.2.13,<4" }, { name = "fastapi", specifier = ">=0.141.1,<1" }, { name = "giteapy", specifier = "==1.0.8" }, @@ -2634,10 +2572,11 @@ requires-dist = [ { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "extra == 'otel-grpc'", specifier = "==1.43.0" }, { name = "opentelemetry-exporter-otlp-proto-http", specifier = "==1.43.0" }, { name = "opentelemetry-sdk", specifier = "==1.43.0" }, + { name = "prometheus-client", specifier = "==0.21.1" }, { name = "protobuf", specifier = "==6.33.6" }, { name = "pydantic", specifier = "==2.13.3" }, { name = "pydantic-settings", specifier = "==2.14.2" }, - { name = "pygithub", specifier = "==1.59.*" }, + { name = "pygithub", specifier = ">=2.10,<3" }, { name = "pyjwt", specifier = ">=2.13.0,<3" }, { name = "python-gitlab", specifier = "==8.3.0" }, { name = "pyyaml", specifier = "==6.0.1" }, @@ -2664,8 +2603,7 @@ dev = [ similar-issue = [ { name = "lancedb", specifier = "==0.38.0" }, { name = "pandas" }, - { name = "pinecone-client" }, - { name = "pinecone-datasets", git = "https://github.com/mrT23/pinecone-datasets.git?rev=main" }, + { name = "pinecone", specifier = "==10.0.0" }, { name = "qdrant-client", specifier = "==1.15.1" }, ] @@ -2685,6 +2623,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/80/6e/4b28b62ecb6aae56769c34a8ff1d661473ec1e9519e2d5f8b2c150086b26/pre_commit-4.6.0-py2.py3-none-any.whl", hash = "sha256:e2cf246f7299edcabcf15f9b0571fdce06058527f0a06535068a86d38089f29b", size = 226472, upload-time = "2026-04-21T20:31:40.092Z" }, ] +[[package]] +name = "prometheus-client" +version = "0.21.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/62/14/7d0f567991f3a9af8d1cd4f619040c93b68f09a02b6d0b6ab1b2d1ded5fe/prometheus_client-0.21.1.tar.gz", hash = "sha256:252505a722ac04b0456be05c05f75f45d760c2911ffc45f2a06bcaed9f3ae3fb", size = 78551, upload-time = "2024-12-03T14:59:12.164Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ff/c2/ab7d37426c179ceb9aeb109a85cda8948bb269b7561a0be870cc656eefe4/prometheus_client-0.21.1-py3-none-any.whl", hash = "sha256:594b45c410d6f4f8888940fe80b5cc2521b305a1fafe1c58609ef715a001f301", size = 54682, upload-time = "2024-12-03T14:59:10.935Z" }, +] + [[package]] name = "propcache" version = "0.5.2" @@ -2994,17 +2941,18 @@ wheels = [ [[package]] name = "pygithub" -version = "1.59.1" +version = "2.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "deprecated" }, { name = "pyjwt", extra = ["crypto"] }, { name = "pynacl" }, { name = "requests" }, + { name = "typing-extensions" }, + { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fb/30/203d3420960853e399de3b85d6613cea1cf17c1cf7fc9716f7ee7e17e0fc/PyGithub-1.59.1.tar.gz", hash = "sha256:c44e3a121c15bf9d3a5cc98d94c9a047a5132a9b01d22264627f58ade9ddc217", size = 3295328, upload-time = "2023-08-03T09:43:01.794Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e1/9b/195603d5371861005a3467c5e4afd02fd0698795a2aa36dc41498b9d879d/pygithub-2.10.0.tar.gz", hash = "sha256:90ff24ef1cd1bd57124c2a3869cafee9d7b066909129ecdaba2c2d1903bc118d", size = 2750952, upload-time = "2026-08-20T10:05:08.327Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2c/71/aff5465d9e3d448a5d4beab1dc7c8dec72037e3ae7e0d856ee08538dc934/PyGithub-1.59.1-py3-none-any.whl", hash = "sha256:3d87a822e6c868142f0c2c4bf16cce4696b5a7a4d142a7bd160e1bdf75bc54a9", size = 342171, upload-time = "2023-08-03T09:43:00.046Z" }, + { url = "https://files.pythonhosted.org/packages/91/71/f314841697a1d52af3e1ea7c5e1c3f09685b64ae4c58ff16b605a866255d/pygithub-2.10.0-py3-none-any.whl", hash = "sha256:192ada2a76e4afc7d6b37e500c9bfeba1731e6506697445a5ba1c4af8bf0b924", size = 455554, upload-time = "2026-08-20T10:05:06.991Z" }, ] [[package]] @@ -3496,20 +3444,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/73/e3/7df5a396e445b9ba49ce9a9437439a4d80042c61c0ade199abf8d16de1ac/ruff-0.16.5-py3-none-win_arm64.whl", hash = "sha256:a64abe90968719b851bb7cedffaa8753fbdbdadab483089682db623f3edc587e", size = 10391564, upload-time = "2026-08-27T16:34:16.064Z" }, ] -[[package]] -name = "s3fs" -version = "0.4.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "botocore" }, - { name = "fsspec", version = "2025.12.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, - { name = "fsspec", version = "2026.6.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d9/9a/504cb277632c4d325beabbd03bb43778f0decb9be22d9e0e6c62f44540c7/s3fs-0.4.2.tar.gz", hash = "sha256:2ca5de8dc18ad7ad350c0bd01aef0406aa5d0fff78a561f0f710f9d9858abdd0", size = 57527, upload-time = "2020-03-31T15:24:26.388Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b8/e4/b8fc59248399d2482b39340ec9be4bb2493846ac23641b43115a7e5cd675/s3fs-0.4.2-py3-none-any.whl", hash = "sha256:91c1dfb45e5217bd441a7a560946fe865ced6225ff7eb0fb459fe6e601a95ed3", size = 19791, upload-time = "2020-03-31T15:24:24.952Z" }, -] - [[package]] name = "s3transfer" version = "0.19.2"