chore: Add support for detecting pi agent and additional cursor agent logic (#4)

Author: blainekasten

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## 0.2.0
+
+- Improve detection for cursor agents
+- Add support for pi coding agent
+
 ## 0.1.1
 
 - Improve detection for cursor usage to avoid mistaking humans as agents

detect_agent/__init__.py

@@ -17,6 +17,7 @@
 OPENCODE: Literal["opencode"] = "opencode"
 GITHUB_COPILOT: Literal["github-copilot"] = "github-copilot"
 GITHUB_COPILOT_CLI: Literal["github-copilot-cli"] = "github-copilot-cli"
+PI: Literal["pi"] = "pi"
 
 KnownAgentNames = Literal[
     "cursor",
@@ -31,6 +32,7 @@
     "augment-cli",
     "opencode",
     "github-copilot",
+    "pi",
 ]
 
 
@@ -51,6 +53,7 @@ class AgentResultNone(TypedDict):
 AgentResult = Union[AgentResultAgent, AgentResultNone]
 
 KNOWN_AGENTS = {
+    "PI": PI,
     "CURSOR": CURSOR,
     "CURSOR_CLI": CURSOR_CLI,
     "CLAUDE": CLAUDE,
@@ -75,10 +78,16 @@ def determine_agent() -> AgentResult:
                 return {"is_agent": True, "agent": {"name": GITHUB_COPILOT}}
             return {"is_agent": True, "agent": {"name": name}}  # type: ignore[return-value]
 
+    if os.environ.get("PI_CODING_AGENT"):
+        return {"is_agent": True, "agent": {"name": PI}}
+
     if os.environ.get("CURSOR_AGENT"):
         return {"is_agent": True, "agent": {"name": CURSOR}}
 
-    if os.environ.get("CURSOR_INVOKED_AS") == "agent":
+    if (
+        os.environ.get("CURSOR_INVOKED_AS") == "agent"
+        or os.environ.get("CURSOR_EXTENSION_HOST_ROLE") == "agent-exec"
+    ):
         return {"is_agent": True, "agent": {"name": CURSOR_CLI}}
 
     if os.environ.get("GEMINI_CLI"):

pyproject.toml

@@ -7,7 +7,7 @@ packages = ["detect_agent"]
 
 [project]
 name = "detect_agent"
-version = "0.1.1"
+version = "0.2.0"
 description = "Detect if code is running in an AI agent or automated development environment"
 readme = "README.md"
 requires-python = ">=3.9"

tests/test_detect_agent.py

@@ -14,8 +14,9 @@
 # Env vars we reset so tests don't leak into each other
 _AGENT_ENV_VARS = (
     "AI_AGENT",
-    "CURSOR_TRACE_ID",
+    "PI_CODING_AGENT",
     "CURSOR_AGENT",
+    "CURSOR_EXTENSION_HOST_ROLE",
     "GEMINI_CLI",
     "CODEX_SANDBOX",
     "CODEX_CI",
@@ -58,6 +59,11 @@ def test_ai_agent_set_detects_custom_agent(self, monkeypatch):
 class TestGitHubCopilotDetection:
     """GitHub Copilot detection."""
 
+    def test_pi_coding_agent_set_detects_pi(self, monkeypatch):
+        monkeypatch.setenv("PI_CODING_AGENT", "1")
+        result = determine_agent()
+        assert result == {"is_agent": True, "agent": {"name": KNOWN_AGENTS["PI"]}}
+
     def test_from_ai_agent_github_copilot(self, monkeypatch):
         monkeypatch.setenv("AI_AGENT", "github-copilot")
         result = determine_agent()
@@ -87,12 +93,8 @@ def test_from_copilot_github_token(self, monkeypatch):
 class TestCursorDetection:
     """Cursor detection."""
 
-    def test_cursor_trace_id_not_set_returns_no_agent(self):
-        result = determine_agent()
-        assert result == {"is_agent": False, "agent": None}
-
-    def test_cursor_trace_id_set_detects_cursor(self, monkeypatch):
-        monkeypatch.setenv("CURSOR_TRACE_ID", "some-uuid")
+    def test_cursor_agent_set_detects_cursor(self, monkeypatch):
+        monkeypatch.setenv("CURSOR_AGENT", "1")
         result = determine_agent()
         assert result == {"is_agent": True, "agent": {"name": KNOWN_AGENTS["CURSOR"]}}
 
@@ -107,8 +109,18 @@ def test_cursor_agent_not_set_returns_no_agent(self):
     def test_cursor_agent_set_detects_cursor_cli(self, monkeypatch):
         monkeypatch.setenv("CURSOR_AGENT", "1")
         result = determine_agent()
+        assert result == {"is_agent": True, "agent": {"name": KNOWN_AGENTS["CURSOR"]}}
+
+    def test_cursor_extension_host_role_agent_exec_detects_cursor_cli(self, monkeypatch):
+        monkeypatch.setenv("CURSOR_EXTENSION_HOST_ROLE", "agent-exec")
+        result = determine_agent()
         assert result == {"is_agent": True, "agent": {"name": KNOWN_AGENTS["CURSOR_CLI"]}}
 
+    def test_cursor_extension_host_role_other_value_returns_no_agent(self, monkeypatch):
+        monkeypatch.setenv("CURSOR_EXTENSION_HOST_ROLE", "something-else")
+        result = determine_agent()
+        assert result == {"is_agent": False, "agent": None}
+
 
 class TestGeminiDetection:
     """Gemini detection."""
@@ -263,7 +275,6 @@ class TestPriorityOrderDetection:
 
     def test_ai_agent_takes_highest_priority(self, monkeypatch):
         monkeypatch.setenv("AI_AGENT", "custom-priority")
-        monkeypatch.setenv("CURSOR_TRACE_ID", "some-uuid")
         monkeypatch.setenv("CURSOR_AGENT", "1")
         monkeypatch.setenv("GEMINI_CLI", "1")
         monkeypatch.setenv("CODEX_SANDBOX", "seatbelt")
@@ -280,24 +291,6 @@ def test_ai_agent_takes_highest_priority(self, monkeypatch):
             result = determine_agent()
         assert result == {"is_agent": True, "agent": {"name": "custom-priority"}}
 
-    def test_cursor_trace_id_takes_priority_over_other_agents(self, monkeypatch):
-        monkeypatch.setenv("CURSOR_TRACE_ID", "some-uuid")
-        monkeypatch.setenv("CURSOR_AGENT", "1")
-        monkeypatch.setenv("GEMINI_CLI", "1")
-        monkeypatch.setenv("CODEX_SANDBOX", "seatbelt")
-        monkeypatch.setenv("ANTIGRAVITY_AGENT", "1")
-        monkeypatch.setenv("AUGMENT_AGENT", "1")
-        monkeypatch.setenv("OPENCODE_CLIENT", "opencode")
-        monkeypatch.setenv("CLAUDE_CODE", "1")
-        monkeypatch.setenv("REPL_ID", "1")
-        monkeypatch.setenv("COPILOT_MODEL", "gpt-5")
-        monkeypatch.setenv("COPILOT_ALLOW_ALL", "true")
-        monkeypatch.setenv("COPILOT_GITHUB_TOKEN", "ghp_xxx")
-        with patch.object(Path, "exists") as mock_exists:
-            mock_exists.side_effect = lambda self: str(self) == DEVIN_LOCAL_PATH
-            result = determine_agent()
-        assert result == {"is_agent": True, "agent": {"name": KNOWN_AGENTS["CURSOR"]}}
-
     def test_cursor_agent_takes_priority_over_remaining_agents(self, monkeypatch):
         monkeypatch.setenv("CURSOR_AGENT", "1")
         monkeypatch.setenv("GEMINI_CLI", "1")
@@ -313,15 +306,14 @@ def test_cursor_agent_takes_priority_over_remaining_agents(self, monkeypatch):
         with patch.object(Path, "exists") as mock_exists:
             mock_exists.side_effect = lambda self: str(self) == DEVIN_LOCAL_PATH
             result = determine_agent()
-        assert result == {"is_agent": True, "agent": {"name": KNOWN_AGENTS["CURSOR_CLI"]}}
+        assert result == {"is_agent": True, "agent": {"name": KNOWN_AGENTS["CURSOR"]}}
 
 
 class TestEdgeCases:
     """Edge cases."""
 
     def test_empty_string_env_vars(self, monkeypatch):
         monkeypatch.setenv("AI_AGENT", "")
-        monkeypatch.setenv("CURSOR_TRACE_ID", "")
         result = determine_agent()
         assert result == {"is_agent": False, "agent": None}
 
@@ -355,7 +347,7 @@ def test_is_agent_boolean(self, monkeypatch):
         assert result["is_agent"] is True
 
     def test_agent_details_when_detected(self, monkeypatch):
-        monkeypatch.setenv("CURSOR_TRACE_ID", "some-id")
+        monkeypatch.setenv("CURSOR_AGENT", "1")
         result = determine_agent()
         assert result["is_agent"] is True
         assert result.get("agent") is not None