fix(integrations): address AgentCOGS PR review feedback

- Remove in-repo docs (belongs in BerriAI/litellm-docs)
- Use litellm_call_id for ingest run_id idempotency
- Lago-style AGENTCOGS_CHARGE_BY proxy attribution; ignore client
  metadata on proxy requests
- Log callback HTTP errors at warning level

Co-authored-by: Cursor <cursoragent@cursor.com>

docs/my-website/docs/observability/agentcogs.md

@@ -1,83 +0,0 @@
-# AgentCOGS - Per-customer margin
-
-[AgentCOGS](https://github.com/vaibhav11123/agentcogs) tracks per-customer LLM cost and gross margin for B2B SaaS (cost + revenue), alongside your existing proxy and observability stack.
-
-## Quick Start
-
-Use one line to send successful completion cost to AgentCOGS:
-
-Get your AgentCOGS [API key and workspace id](https://github.com/vaibhav11123/agentcogs/blob/main/docs/quickstart.md).
-
-```python
-import os
-import litellm
-
-os.environ["AGENTCOGS_API_KEY"] = ""
-os.environ["AGENTCOGS_WORKSPACE_ID"] = ""
-# optional — defaults to https://api.agentcogs.dev
-os.environ["AGENTCOGS_ENDPOINT"] = "https://api.agentcogs.dev"
-
-litellm.success_callback = ["agentcogs"]
-
-response = litellm.completion(
-    model="gpt-4o-mini",
-    messages=[{"role": "user", "content": "Hi"}],
-    user="your_customer_id",  # 👈 B2B tenant id → AgentCOGS customer_id
-    metadata={"agentcogs_workflow_id": "support_bot"},
-)
-```
-
-- SDK
-- PROXY
-
-```yaml
-litellm_settings:
-  callbacks: ["agentcogs"]
-```
-
-Set `AGENTCOGS_API_KEY`, `AGENTCOGS_WORKSPACE_ID`, and optionally `AGENTCOGS_ENDPOINT` in the proxy environment.
-
-## Environment variables
-
-| Variable | Required | Description |
-|----------|----------|-------------|
-| `AGENTCOGS_API_KEY` | Yes | Workspace API key (`acg_live_...`) |
-| `AGENTCOGS_WORKSPACE_ID` | Yes | Workspace UUID |
-| `AGENTCOGS_ENDPOINT` | No | API base URL (default `https://api.agentcogs.dev`) |
-
-## Tenant attribution
-
-Pass `user=` on each completion (same pattern as [Lago](./lago.md)). LiteLLM maps it to AgentCOGS `customer_id`.
-
-Alternatively set `metadata.agentcogs_customer_id` if you cannot use `user`.
-
-Completions without a customer id are skipped (proxy traffic is not blocked).
-
-## What LiteLLM sends
-
-Each successful or failed completion POSTs to `POST /v1/ingest`:
-
-```json
-{
-  "run_id": "<uuid>",
-  "workspace_id": "<AGENTCOGS_WORKSPACE_ID>",
-  "customer_id": "<user or metadata.agentcogs_customer_id>",
-  "workflow_id": "default",
-  "ts": 1710000000,
-  "status": "completed",
-  "total_usd": 0.0012,
-  "models": {
-    "gpt-4o-mini": {
-      "input_tokens": 10,
-      "output_tokens": 5,
-      "usd": 0.0012
-    }
-  },
-  "metadata": { "source": "litellm" }
-}
-```
-
-## Learn more
-
-- [AgentCOGS quickstart](https://github.com/vaibhav11123/agentcogs/blob/main/docs/quickstart.md)
-- [LiteLLM callback (user-landed)](https://github.com/vaibhav11123/agentcogs/blob/main/docs/integrations/litellm.md)

litellm/integrations/agentcogs.py

@@ -3,7 +3,7 @@
 import json
 import os
 import time
-from typing import Any, Optional
+from typing import Any, Literal, Optional
 
 import litellm
 from litellm._logging import verbose_logger
@@ -57,6 +57,50 @@ def _metadata_from_kwargs(kwargs: dict) -> dict:
     return meta if isinstance(meta, dict) else {}
 
 
+def _resolve_customer_id(kwargs: dict) -> Optional[str]:
+    """
+    Resolve tenant id for AgentCOGS ingest.
+
+    Proxy mode (default): same attribution sources as Lago — end_user_id from proxy
+    body, or authenticated key metadata (user_id / team_id). Client-supplied
+    metadata.agentcogs_customer_id is not trusted on proxy requests.
+
+    Direct SDK mode (no proxy_server_request): kwargs user or metadata.agentcogs_customer_id.
+    """
+    litellm_params = kwargs.get("litellm_params", {}) or {}
+    meta = litellm_params.get("metadata", {}) or {}
+    proxy_server_request = litellm_params.get("proxy_server_request")
+    is_proxy = bool(proxy_server_request)
+
+    proxy_body = (proxy_server_request or {}).get("body") or {}
+    end_user_id = proxy_body.get("user") or kwargs.get("user")
+    user_id = meta.get("user_api_key_user_id")
+    team_id = meta.get("user_api_key_team_id")
+
+    charge_by: Literal["end_user_id", "team_id", "user_id"] = "end_user_id"
+    if os.getenv("AGENTCOGS_CHARGE_BY") is not None and isinstance(
+        os.environ["AGENTCOGS_CHARGE_BY"], str
+    ):
+        if os.environ["AGENTCOGS_CHARGE_BY"] in ("end_user_id", "user_id", "team_id"):
+            charge_by = os.environ["AGENTCOGS_CHARGE_BY"]  # type: ignore
+        else:
+            raise Exception("invalid AGENTCOGS_CHARGE_BY set")
+
+    if charge_by == "end_user_id":
+        customer_id = end_user_id
+    elif charge_by == "team_id":
+        customer_id = team_id
+    else:
+        customer_id = user_id
+
+    if customer_id is None and not is_proxy:
+        customer_id = _metadata_from_kwargs(kwargs).get("agentcogs_customer_id")
+
+    if customer_id is None:
+        return None
+    return str(customer_id)
+
+
 class AgentCOGSLogger(CustomLogger):
     """POST per-completion cost to AgentCOGS /v1/ingest for B2B per-customer margin."""
 
@@ -97,11 +141,11 @@ class AgentCOGSLogger(CustomLogger):
         start_time: Any = None,
         error: Optional[str] = None,
     ) -> Optional[dict]:
-        meta = _metadata_from_kwargs(kwargs)
-        customer_id = kwargs.get("user") or meta.get("agentcogs_customer_id")
+        customer_id = _resolve_customer_id(kwargs)
         if not customer_id:
             return None
 
+        meta = _metadata_from_kwargs(kwargs)
         model = kwargs.get("model") or "unknown"
         cost = float(kwargs.get("response_cost") or 0)
         usage = _extract_usage(kwargs, response_obj)
@@ -113,9 +157,9 @@ class AgentCOGSLogger(CustomLogger):
             ts = int(time.time())
 
         return {
-            "run_id": str(uuid.uuid4()),
+            "run_id": str(kwargs.get("litellm_call_id") or uuid.uuid4()),
             "workspace_id": os.environ["AGENTCOGS_WORKSPACE_ID"],
-            "customer_id": str(customer_id),
+            "customer_id": customer_id,
             "workflow_id": meta.get("agentcogs_workflow_id", "default"),
             "ts": ts,
             "status": status,
@@ -141,7 +185,7 @@ class AgentCOGSLogger(CustomLogger):
             )
             response.raise_for_status()
         except Exception as e:
-            verbose_logger.debug("AgentCOGS callback error: {}".format(e))
+            verbose_logger.warning("AgentCOGS callback error: {}".format(e))
 
     def _sync_post(self, event: dict) -> None:
         try:
@@ -152,7 +196,7 @@ class AgentCOGSLogger(CustomLogger):
             )
             response.raise_for_status()
         except Exception as e:
-            verbose_logger.debug("AgentCOGS callback error: {}".format(e))
+            verbose_logger.warning("AgentCOGS callback error: {}".format(e))
 
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
         event = self._build_event(

tests/test_litellm/integrations/test_agentcogs.py

@@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-from litellm.integrations.agentcogs import AgentCOGSLogger
+from litellm.integrations.agentcogs import AgentCOGSLogger, _resolve_customer_id
 
 
 class TestAgentCOGSIntegration:
@@ -18,6 +18,7 @@ class TestAgentCOGSIntegration:
             "AGENTCOGS_API_KEY",
             "AGENTCOGS_WORKSPACE_ID",
             "AGENTCOGS_ENDPOINT",
+            "AGENTCOGS_CHARGE_BY",
         ):
             os.environ.pop(key, None)
 
@@ -58,7 +59,8 @@ class TestAgentCOGSIntegration:
         }
         assert logger._build_event(kwargs, {}, status="completed") is None
 
-    def test_build_event_customer_from_metadata(self):
+    def test_build_event_customer_from_metadata_sdk_mode(self):
+        """SDK/direct calls (no proxy): metadata.agentcogs_customer_id is allowed."""
         logger = AgentCOGSLogger()
         kwargs = {
             "litellm_params": {
@@ -71,6 +73,49 @@ class TestAgentCOGSIntegration:
         assert event is not None
         assert event["customer_id"] == "meta_tenant"
 
+    def test_run_id_uses_litellm_call_id(self):
+        logger = AgentCOGSLogger()
+        kwargs = {
+            "user": "acme",
+            "litellm_call_id": "call-abc-123",
+            "model": "gpt-4",
+            "response_cost": 0.001,
+        }
+        event = logger._build_event(kwargs, {}, status="completed")
+        assert event is not None
+        assert event["run_id"] == "call-abc-123"
+
+    def test_proxy_uses_end_user_from_body(self):
+        kwargs = {
+            "litellm_params": {
+                "proxy_server_request": {"body": {"user": "proxy_tenant"}},
+                "metadata": {
+                    "agentcogs_customer_id": "malicious_tenant",
+                    "user_api_key_user_id": "key_user",
+                },
+            },
+        }
+        assert _resolve_customer_id(kwargs) == "proxy_tenant"
+
+    def test_proxy_ignores_client_metadata_customer_id(self):
+        kwargs = {
+            "litellm_params": {
+                "proxy_server_request": {"body": {}},
+                "metadata": {"agentcogs_customer_id": "malicious_tenant"},
+            },
+        }
+        assert _resolve_customer_id(kwargs) is None
+
+    def test_proxy_charge_by_team_id(self):
+        os.environ["AGENTCOGS_CHARGE_BY"] = "team_id"
+        kwargs = {
+            "litellm_params": {
+                "proxy_server_request": {"body": {"user": "end_user"}},
+                "metadata": {"user_api_key_team_id": "team-99"},
+            },
+        }
+        assert _resolve_customer_id(kwargs) == "team-99"
+
     def test_build_event_error_status(self):
         logger = AgentCOGSLogger()
         kwargs = {