diff --git a/scripts/push_predictions.py b/scripts/push_predictions.py index baa91eb..9ac9712 100644 --- a/scripts/push_predictions.py +++ b/scripts/push_predictions.py @@ -15,6 +15,10 @@ export LABEL_STUDIO_API_KEY=xxxx +新版 Label Studio 默认只认 **JWT 个人访问令牌**(组织设置里 `legacy_api_tokens_enabled=false`)。 +脚本会自动识别:令牌形如 JWT(两个点分三段)时先用 `/api/token/refresh/` 换取 access token +并以 `Bearer` 发送,否则按旧式 `Token` 发送。 + 图片路径从任务的 `/data/local-files/?d=...` 还原,所以 `--document-root` 必须和 Label Studio 的 `LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT` 一致。 @@ -49,7 +53,25 @@ class LabelStudio: self.url = url.rstrip("/") self.session = requests.Session() - self.session.headers["Authorization"] = f"Token {token}" + self.session.headers["Authorization"] = self._authorization(token) + + def _authorization(self, token: str) -> str: + """区分两种令牌。 + + 新版默认关闭了旧式 `Token xxx`(组织设置 legacy_api_tokens_enabled=false), + 个人访问令牌是一个 JWT refresh token,要先换成 access token 再以 Bearer 发送。 + """ + if token.count(".") != 2: # 不是 JWT,按旧式 Token 处理 + return f"Token {token}" + response = self.session.post( + f"{self.url}/api/token/refresh/", json={"refresh": token}, timeout=60 + ) + if response.status_code >= 400: + raise RuntimeError( + f"用个人访问令牌换取 access token 失败(HTTP {response.status_code}):" + f"{response.text[:200]}" + ) + return f"Bearer {response.json()['access']}" def tasks(self, project: int, page_size: int = 500) -> list[dict]: """列出项目里的全部任务(只取 id 与 data,分页拉完)。""" @@ -179,7 +201,11 @@ def main() -> int: logger.error("%s", exc) return 1 - client = LabelStudio(args.url, args.token) + try: + client = LabelStudio(args.url, args.token) + except RuntimeError as exc: + logger.error("%s", exc) + return 1 logger.info("拉取项目 %d 的任务列表…", args.project) tasks = client.tasks(args.project) if args.limit: diff --git a/tests/test_push_predictions.py b/tests/test_push_predictions.py index 17310d4..ca24ebe 100644 --- a/tests/test_push_predictions.py +++ b/tests/test_push_predictions.py @@ -33,6 +33,10 @@ class _FakeResponse: self._payload = payload self.status_code = status + @property + def text(self): + return str(self._payload) + def json(self): return self._payload @@ -86,3 +90,42 @@ def test_tasks_stops_on_empty_page(client): client.session = _FakeSession([{"tasks": [{"id": 1}]}, {"tasks": []}]) assert [t["id"] for t in client.tasks(39)] == [1] assert len(client.session.calls) == 2 + + +# ---------------- 两种令牌 ---------------- + + +class _TokenSession: + def __init__(self, access="ACCESS", status=200): + self.headers = {} + self.access, self.status, self.posted = access, status, [] + + def post(self, url, json=None, timeout=None): + self.posted.append((url, json)) + return _FakeResponse({"access": self.access}, self.status) + + +def test_legacy_token_is_sent_as_token(monkeypatch): + client = pp.LabelStudio.__new__(pp.LabelStudio) + client.url = "http://ls" + client.session = _TokenSession() + assert client._authorization("c657261234567890") == "Token c657261234567890" + assert client.session.posted == [] # 旧式令牌不需要换取 + + +def test_jwt_personal_token_is_exchanged_for_a_bearer(monkeypatch): + """新版默认禁用旧式 token,个人访问令牌要先换成 access token。""" + client = pp.LabelStudio.__new__(pp.LabelStudio) + client.url = "http://ls" + client.session = _TokenSession(access="AAA") + assert client._authorization("aaa.bbb.ccc") == "Bearer AAA" + assert client.session.posted[0][0] == "http://ls/api/token/refresh/" + assert client.session.posted[0][1] == {"refresh": "aaa.bbb.ccc"} + + +def test_failed_exchange_reports_clearly(): + client = pp.LabelStudio.__new__(pp.LabelStudio) + client.url = "http://ls" + client.session = _TokenSession(status=401) + with pytest.raises(RuntimeError, match="access token 失败"): + client._authorization("aaa.bbb.ccc")