_detect_agent.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. # Copyright 2026 The HuggingFace Team. All rights reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """Detect whether the process is being invoked by an AI coding agent.
  15. Detection is based on environment variables that AI agents set in their shell
  16. sessions. ``AI_AGENT`` and ``AGENT`` are treated as a universal standard (any
  17. tool can set it); the remaining checks are tool-specific and ordered by
  18. prevalence.
  19. Inspired by ``@vercel/detect-agent`` (https://github.com/vercel/vercel/tree/main/packages/detect-agent).
  20. """
  21. import os
  22. from typing import Optional
  23. # Standard env vars — value is used as the agent name directly.
  24. _STANDARD_AGENT_VARS: tuple[str, ...] = ("AI_AGENT", "AGENT")
  25. # NOTE: ``cowork`` must appear before ``claude-code`` so the more specific
  26. # signal takes priority when both ``CLAUDE_CODE`` and ``CLAUDE_CODE_IS_COWORK``
  27. # are set.
  28. _TOOL_AGENTS: tuple[tuple[tuple[str, ...], str], ...] = (
  29. (("ANTIGRAVITY_AGENT",), "antigravity"),
  30. (("AUGMENT_AGENT",), "augment-cli"),
  31. (("CLINE_ACTIVE",), "cline"),
  32. (("CLAUDE_CODE_IS_COWORK",), "cowork"),
  33. (("CLAUDECODE", "CLAUDE_CODE"), "claude-code"),
  34. (("CODEX_SANDBOX", "CODEX_CI", "CODEX_THREAD_ID"), "codex"),
  35. (("CURSOR_TRACE_ID",), "cursor"),
  36. (("CURSOR_AGENT",), "cursor-cli"),
  37. (("GEMINI_CLI",), "gemini"),
  38. (("COPILOT_MODEL", "COPILOT_ALLOW_ALL", "COPILOT_GITHUB_TOKEN"), "github-copilot"),
  39. (("GOOSE_TERMINAL",), "goose"),
  40. (("OPENCLAW_SHELL",), "openclaw"),
  41. (("OPENCODE_CLIENT",), "opencode"),
  42. (("REPL_ID",), "replit"),
  43. (("ROO_ACTIVE",), "roo-code"),
  44. (("TRAE_AI_SHELL_ID",), "trae"),
  45. )
  46. _KNOWN_AGENTS = {"devin"} | {agent for _, agent in _TOOL_AGENTS}
  47. def detect_agent() -> Optional[str]:
  48. """Return the name of the detected AI agent or ``None``.
  49. Checks environment variables in priority order and returns on the first
  50. match. When ``AI_AGENT`` or ``AGENT`` is set, the value is checked against
  51. known agent names, unrecognized values are returned as ``"unknown"``.
  52. """
  53. for var in _STANDARD_AGENT_VARS:
  54. name = os.environ.get(var, "").strip().lower()
  55. if name:
  56. return name if name in _KNOWN_AGENTS else "unknown"
  57. for env_vars, agent_name in _TOOL_AGENTS:
  58. if any(os.environ.get(var) for var in env_vars):
  59. return agent_name
  60. return None
  61. def is_agent() -> bool:
  62. """Return ``True`` if the process is being invoked by an AI coding agent."""
  63. return detect_agent() is not None