ap-client 0.4.0__tar.gz → 0.5.0.dev0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: ap-client
3
- Version: 0.4.0
3
+ Version: 0.5.0.dev0
4
4
  Summary: Agent Platform API Client & CLI
5
5
  Requires-Python: >=3.10
6
6
  Requires-Dist: pyyaml>=6.0
@@ -9,9 +9,9 @@ Requires-Dist: rich>=13.0.0
9
9
  Requires-Dist: typer>=0.9.0
10
10
  Requires-Dist: websockets>=13.0
11
11
  Provides-Extra: all
12
- Requires-Dist: instance-repo[oss]<2.0.0,>=1.1.2; extra == 'all'
12
+ Requires-Dist: instance-repo[oss]<2.0.0,>=1.1.4; extra == 'all'
13
13
  Provides-Extra: dataset
14
- Requires-Dist: instance-repo[oss]<2.0.0,>=1.1.2; extra == 'dataset'
14
+ Requires-Dist: instance-repo[oss]<2.0.0,>=1.1.4; extra == 'dataset'
15
15
  Description-Content-Type: text/markdown
16
16
 
17
17
  A lightweight Python SDK and command line interface for Agent Platform. It provides helpers for configuring API access and managing templates, datasets, jobs, and groups.
@@ -0,0 +1,181 @@
1
+ """``ap akpool`` 命令族 —— Account Pool(托管 AK)观测操作。
2
+
3
+ * ``akpool rpm get`` → apiserver ``GET /apis/v1/model-api-keys/rpm-usage``
4
+ (中心侧,走 ``APIClient._get_central``),按 AK + 模型查询当前 RPM 用量与上限。
5
+
6
+ AK 的两种指定方式互斥且必填其一:``--ak-id``(uak-xxx,仅平台托管 AK)或
7
+ ``--masked-ak``(掩码后 6 位,未托管 AK 走这条路)。服务端同口径(都缺或都给 →
8
+ 400),CLI 侧先拦一道,不发请求。``--time-range`` 是 Go duration 风格
9
+ (``30s``/``5m``/``1h``);显式 ``--from``/``--to``(RFC3339)时覆盖
10
+ ``--time-range``。
11
+
12
+ 响应是嵌套结构(无信封):``usage.current_rpm``/``usage.sampled_at``、
13
+ ``rpm.limit``/``rpm.limit_state``/``rpm.synced_at``、``usage_status``、
14
+ ``time_range``、``ak_id``/``masked_api_key``/``model_id``。json/yaml 原样透传;
15
+ plain/table 摊平成单行 key=value 展示。
16
+ """
17
+
18
+ from __future__ import annotations
19
+
20
+ import re
21
+ from typing import Any, Optional
22
+
23
+ import typer
24
+ from ap_client import get_client
25
+
26
+ __all__ = ["akpool_app", "register"]
27
+
28
+ akpool_app = typer.Typer(help="Account Pool operations (managed API keys)")
29
+ rpm_app = typer.Typer(help="Account Pool RPM usage operations")
30
+
31
+ akpool_app.add_typer(rpm_app, name="rpm")
32
+
33
+
34
+ def register(app: typer.Typer) -> None:
35
+ """把 ``akpool`` 挂到 ``ap`` 根命令上。"""
36
+ app.add_typer(akpool_app, name="akpool")
37
+
38
+
39
+ _FORMAT_HELP = "Output format: plain/table/json/yaml (default: AP_FORMAT or command default)"
40
+ _TIME_RANGE_HELP = "Lookback window for the RPM reading (Go duration, e.g. 30s/5m/1h)"
41
+ _ISO_TIME_HELP = "RFC3339 timestamp, e.g. 2026-09-21T12:00:00+08:00"
42
+
43
+ _TIME_RANGE_RE = re.compile(r"^\d+[smh]$")
44
+
45
+ #: plain/table 输出的字段顺序,取值自嵌套响应(见 _flatten_rpm)。
46
+ _RPM_KEYS = (
47
+ "current_rpm",
48
+ "rpm_limit",
49
+ "model",
50
+ "sampled_at",
51
+ "time_range",
52
+ "ak_id",
53
+ "usage_status",
54
+ )
55
+
56
+
57
+ def _fmt(output_format: Optional[str]) -> str:
58
+ from .cli import _normalize_output_format
59
+
60
+ return _normalize_output_format(output_format, keep_table=True)
61
+
62
+
63
+ def _check_time_range(value: str) -> str:
64
+ normalized = str(value or "").strip().lower()
65
+ if not _TIME_RANGE_RE.fullmatch(normalized):
66
+ raise typer.BadParameter(
67
+ f"--time-range must be a duration like 30s/5m/1h, got {value!r}"
68
+ )
69
+ return normalized
70
+
71
+
72
+ def _flatten_rpm(payload: dict) -> dict:
73
+ """把嵌套的 rpm-usage 响应摊平成 plain/table 展示用的扁平 dict。
74
+
75
+ null 值渲染为 ``-``(如未托管 AK 的 ``ak_id``、未知限额的 ``rpm.limit``)。
76
+ """
77
+ rpm = payload.get("rpm")
78
+ usage = payload.get("usage")
79
+ flat = {
80
+ "current_rpm": usage.get("current_rpm") if isinstance(usage, dict) else None,
81
+ "rpm_limit": rpm.get("limit") if isinstance(rpm, dict) else None,
82
+ "model": payload.get("model_id"),
83
+ "sampled_at": usage.get("sampled_at") if isinstance(usage, dict) else None,
84
+ "time_range": payload.get("time_range"),
85
+ "ak_id": payload.get("ak_id"),
86
+ "usage_status": payload.get("usage_status"),
87
+ }
88
+ return {key: (value if value is not None else "-") for key, value in flat.items()}
89
+
90
+
91
+ def _emit_rpm(payload: Any, output_format: str) -> None:
92
+ from .cli import _print_formatted, _print_rich_table
93
+
94
+ if not isinstance(payload, dict):
95
+ payload = {"value": payload}
96
+ if output_format in ("json", "yaml"):
97
+ _print_formatted(payload, output_format)
98
+ return
99
+ flat = _flatten_rpm(payload)
100
+ keys = [key for key in _RPM_KEYS if key in flat]
101
+ keys += [key for key in flat if key not in keys]
102
+ if output_format == "table":
103
+ _print_rich_table([flat], [(key, key) for key in keys])
104
+ return
105
+ _print_rpm_plain(payload)
106
+
107
+
108
+ def _print_rpm_plain(payload: dict) -> None:
109
+ """plain 输出:标题行(usage_status) + 对齐 key/value,与 _print_job_get_plain 同风格。
110
+
111
+ current_rpm 在限额已知时附利用率百分比;rpm_limit 未知/不限时展示 limit_state。
112
+ ak_id / masked_api_key 为 null(未托管 AK)时整行省略。
113
+ """
114
+ from .cli import _print_key_values
115
+
116
+ rpm = payload.get("rpm")
117
+ usage = payload.get("usage")
118
+ rpm = rpm if isinstance(rpm, dict) else {}
119
+ usage = usage if isinstance(usage, dict) else {}
120
+
121
+ current = usage.get("current_rpm")
122
+ limit = rpm.get("limit")
123
+ limit_state = rpm.get("limit_state") or "unknown"
124
+
125
+ if isinstance(current, (int, float)) and isinstance(limit, (int, float)) and limit > 0:
126
+ current_text = f"{current} ({current / limit * 100:.1f}% of limit)"
127
+ else:
128
+ current_text = current
129
+ limit_text = f"{limit} ({limit_state})" if limit is not None else limit_state
130
+
131
+ typer.echo(f"AK RPM Usage ({payload.get('usage_status') or 'unknown'})")
132
+ _print_key_values(
133
+ [
134
+ ("current_rpm", current_text),
135
+ ("rpm_limit", limit_text),
136
+ ("model", payload.get("model_id")),
137
+ ("ak_id", payload.get("ak_id")),
138
+ ("masked_api_key", payload.get("masked_api_key")),
139
+ ("sampled_at", usage.get("sampled_at")),
140
+ ("time_range", payload.get("time_range")),
141
+ ]
142
+ )
143
+
144
+
145
+ @rpm_app.command("get")
146
+ def rpm_get(
147
+ ak_id: Optional[str] = typer.Option(
148
+ None, "--ak-id", help="Managed AK ID (uak-xxx); mutually exclusive with --masked-ak"
149
+ ),
150
+ masked_ak: Optional[str] = typer.Option(
151
+ None, "--masked-ak", help="Last 6 characters of the AK; mutually exclusive with --ak-id"
152
+ ),
153
+ model: str = typer.Option(..., "--model", help="Model the RPM reading applies to"),
154
+ time_range: str = typer.Option("5m", "--time-range", help=_TIME_RANGE_HELP),
155
+ from_time: Optional[str] = typer.Option(
156
+ None, "--from", help=f"Window start ({_ISO_TIME_HELP}); overrides --time-range"
157
+ ),
158
+ to_time: Optional[str] = typer.Option(
159
+ None, "--to", help=f"Window end ({_ISO_TIME_HELP}); overrides --time-range"
160
+ ),
161
+ output_format: str = typer.Option(None, "--format", help=_FORMAT_HELP),
162
+ ):
163
+ """Show the current RPM usage and limit of a managed AK for one model."""
164
+ resolved_ak_id = (ak_id or "").strip()
165
+ resolved_masked_ak = (masked_ak or "").strip()
166
+ if bool(resolved_ak_id) == bool(resolved_masked_ak):
167
+ raise typer.BadParameter("exactly one of --ak-id or --masked-ak is required")
168
+ if not model.strip():
169
+ raise typer.BadParameter("--model must not be empty")
170
+ normalized_range = _check_time_range(time_range)
171
+ output_format = _fmt(output_format)
172
+
173
+ result = get_client().get_akpool_rpm(
174
+ model=model.strip(),
175
+ ak_id=resolved_ak_id or None,
176
+ masked_ak=resolved_masked_ak or None,
177
+ time_range=normalized_range,
178
+ from_time=from_time,
179
+ to_time=to_time,
180
+ )
181
+ _emit_rpm(result, output_format)
@@ -9,10 +9,11 @@ import sys
9
9
  import time
10
10
  import uuid
11
11
  from email.utils import parsedate_to_datetime
12
- from typing import Any, Optional, Union
12
+ from typing import Any, Literal, Optional, Union
13
13
  from urllib.parse import parse_qsl, quote, urlencode, urlsplit, urlunsplit
14
14
 
15
15
  import requests
16
+ from urllib3.util import Timeout as HTTPTimeout
16
17
 
17
18
  from .config import Config, ConfigurationError, get_config, has_cluster_header
18
19
  from .managed_ak import require_managed_ak_ack, validate_ak_id, validate_ak_selection
@@ -126,7 +127,6 @@ EXTRA_AGG_OPS = ("avg", "sum")
126
127
  _OTHERS_MODEL = {
127
128
  "name": "Others",
128
129
  "series_name": "",
129
- "is_default": False,
130
130
  "source_id": "",
131
131
  "series_id": "",
132
132
  }
@@ -171,7 +171,9 @@ def _merge_meta_tags(
171
171
  conflicts = [k for k, v in provided.items() if v and str(v).strip() and k in existing_keys]
172
172
  if conflicts:
173
173
  raise ValueError(
174
- "tags 中已包含 " + ", ".join(conflicts) + ",与 meta 参数冲突,请只保留一种方式"
174
+ "tags already contains "
175
+ + ", ".join(conflicts)
176
+ + "; it conflicts with the meta parameter, keep only one"
175
177
  )
176
178
  extra = [
177
179
  f"{key}:{str(value).strip()}"
@@ -343,6 +345,11 @@ def _split_paged(payload: Any) -> tuple[list[dict], dict]:
343
345
  )
344
346
 
345
347
 
348
+ # apiserver meta-tags 端点单页上限(handler.metaTagPaging 截断到 200;
349
+ # 缺省仅 20,目录查询必须显式带上并翻页)。
350
+ _META_TAG_PAGE_SIZE = 200
351
+
352
+
346
353
  class APIClient:
347
354
  """Agent Platform API client."""
348
355
 
@@ -401,6 +408,8 @@ class APIClient:
401
408
  params: Optional[dict] = None,
402
409
  json_body: Any = None,
403
410
  timeout: TimeoutType = None,
411
+ deadline: Optional[float] = None,
412
+ _return_response: bool = False,
404
413
  ) -> Any:
405
414
  """Send an HTTP request with X-Request-ID and optional verbose logging."""
406
415
  self._require_cluster_config()
@@ -412,6 +421,7 @@ class APIClient:
412
421
  verbose = self.config.verbose
413
422
  max_attempts = self._max_attempts_for(method)
414
423
  for attempt in range(1, max_attempts + 1):
424
+ attempt_timeout = _deadline_timeout(timeout, deadline)
415
425
  request_id = str(uuid.uuid4())
416
426
  headers = {"X-Request-ID": request_id}
417
427
  if verbose:
@@ -435,9 +445,11 @@ class APIClient:
435
445
  params=params,
436
446
  json=json_body,
437
447
  headers=headers,
438
- timeout=timeout,
448
+ timeout=attempt_timeout,
439
449
  )
440
450
  except requests.RequestException as exc:
451
+ if deadline is not None:
452
+ _manifest_budget_remaining(deadline)
441
453
  duration_ms = (time.perf_counter() - start) * 1000
442
454
  if verbose:
443
455
  self._log_exception(method, url, request_id, duration_ms, exc)
@@ -453,10 +465,12 @@ class APIClient:
453
465
  delay,
454
466
  error=f"{type(exc).__name__}: {exc}",
455
467
  )
456
- time.sleep(delay)
468
+ _deadline_sleep(delay, deadline)
457
469
  continue
458
470
  raise APIError(None, str(exc), request_id, None) from exc
459
471
  duration_ms = (time.perf_counter() - start) * 1000
472
+ if deadline is not None:
473
+ _manifest_budget_remaining(deadline)
460
474
  if verbose:
461
475
  self._log_response(method, url, request_id, resp, duration_ms)
462
476
  if not resp.ok:
@@ -472,10 +486,12 @@ class APIClient:
472
486
  delay,
473
487
  status=resp.status_code,
474
488
  )
475
- time.sleep(delay)
489
+ _deadline_sleep(delay, deadline)
476
490
  continue
477
491
  detail = _response_error_detail(resp)
478
492
  raise APIError(resp.status_code, detail, request_id, resp)
493
+ if _return_response:
494
+ return resp, request_id
479
495
  if not resp.content:
480
496
  return None
481
497
  try:
@@ -977,24 +993,142 @@ class APIClient:
977
993
  return item
978
994
  raise APIError(404, f"meta model {name!r} not found", "client-filter", None)
979
995
 
996
+ def _get_meta_tag_catalog(
997
+ self, path: str, params: Optional[dict] = None
998
+ ) -> tuple[list[dict], dict]:
999
+ """按 benchmarks 同口径自动翻页取全量目录记录,返回 ``(records, pagination)``。
1000
+
1001
+ apiserver ``meta-tags`` 端点分页参数为 ``page`` / ``page_size``(缺省
1002
+ 20、上限 200,见 handler.metaTagPaging);单页取用会把值目录静默截断
1003
+ (``meta-job-type`` 漏掉 ``Eval-Eval`` 时 ``default`` 会被算成空串),
1004
+ 故显式传上限并翻页。
1005
+ """
1006
+ records: list[dict] = []
1007
+ pagination: dict = {}
1008
+ page = 1
1009
+ # 与 list_benchmarks 同口径的 100 页上限,防分页异常导致死循环
1010
+ max_pages = 100
1011
+ while page <= max_pages:
1012
+ query = dict(params or {})
1013
+ query["page"] = page
1014
+ query["page_size"] = _META_TAG_PAGE_SIZE
1015
+ page_records, pagination = _split_paged(self._get_central(path, params=query))
1016
+ records.extend(page_records)
1017
+ if not page_records or not pagination:
1018
+ break
1019
+ try:
1020
+ total = int(pagination.get("total"))
1021
+ except (TypeError, ValueError):
1022
+ total = None
1023
+ if total is not None:
1024
+ if len(records) >= total:
1025
+ break
1026
+ elif len(page_records) < _META_TAG_PAGE_SIZE:
1027
+ break
1028
+ page += 1
1029
+ return records, pagination
1030
+
1031
+ def list_meta_tags(self) -> dict:
1032
+ """meta 标签键目录(apiserver GET /apis/v1/meta-tags/keys,返回 {items:[...]})。
1033
+
1034
+ apiserver 用 PagedSuccess envelope({code,data,pagination})返回,
1035
+ 按基准 benchmarks 客户端同口径经 _split_paged 拆包并自动翻页;
1036
+ 裸列表等旧形态兜底为空分页。
1037
+ """
1038
+ records, pagination = self._get_meta_tag_catalog("/meta-tags/keys")
1039
+ result = {"items": records}
1040
+ if pagination:
1041
+ result["pagination"] = pagination
1042
+ return result
1043
+
1044
+ def get_meta_tags(self, key: str) -> dict:
1045
+ """单个 meta 键的值目录(GET /apis/v1/meta-tags?key=,返回 {items:[...]})。
1046
+
1047
+ 与 ``list_meta_tags`` 同样自动翻页,避免值目录被单页默认大小截断。
1048
+ """
1049
+ records, pagination = self._get_meta_tag_catalog("/meta-tags", {"key": key})
1050
+ result = {"items": records}
1051
+ if pagination:
1052
+ result["pagination"] = pagination
1053
+ return result
1054
+
980
1055
  def list_meta_job_types(self) -> dict:
981
- """任务类型词表(本地常量,不走接口)。
1056
+ """任务类型词表(目录驱动:GET /apis/v1/meta-tags?key=meta-job-type)。
982
1057
 
983
- job-type 词表固定为 Eval-Eval(默认)+ Others(兜底手动输入),
984
- 不依赖服务端可用性。default 固定 ``Eval-Eval``(job-type 词表
985
- 不再走集群配置,``AP_META_DEFAULT_JOB_TYPE`` env 已删)。
1058
+ 仅取目录中 ``is_enabled`` 的值(缺失按启用解析)+ 末尾 ``Others``
1059
+ 兜底行;``default`` 为 ``Eval-Eval``(在目录中时,否则空串)。
1060
+ 目录请求失败或为空 → 回退本地常量(Eval-Eval + Others,fail-open),
1061
+ 并在 stderr 打一行告警,词表降级可见但不阻塞命令。
986
1062
  """
987
- default = "Eval-Eval"
1063
+ error: Exception | None = None
1064
+ try:
1065
+ items = self.get_meta_tags("meta-job-type").get("items") or []
1066
+ except Exception as exc: # fail-open:目录不可用不阻塞词表命令
1067
+ error = exc
1068
+ items = []
1069
+ if items:
1070
+ options = []
1071
+ for item in items:
1072
+ value = str(item.get("value") or "").strip()
1073
+ if value and item.get("is_enabled", True):
1074
+ options.append({"value": value})
1075
+ if options:
1076
+ default = "Eval-Eval" if any(o["value"] == "Eval-Eval" for o in options) else ""
1077
+ return {
1078
+ "meta_job_type": {
1079
+ "default": default,
1080
+ "options": [*options, {"value": "Others"}],
1081
+ }
1082
+ }
1083
+ _write_stderr(
1084
+ [
1085
+ "warning: meta-job-type catalog unavailable or empty, "
1086
+ f"falling back to local constants (reason={error or 'empty catalog'})"
1087
+ ]
1088
+ )
988
1089
  return {
989
1090
  "meta_job_type": {
990
- "default": default,
991
- "options": [
992
- {"value": default, "is_default": True},
993
- {"value": "Others", "is_default": False},
994
- ],
1091
+ "default": "Eval-Eval",
1092
+ "options": [{"value": "Eval-Eval"}, {"value": "Others"}],
995
1093
  }
996
1094
  }
997
1095
 
1096
+ # ==================== Account Pool operations ====================
1097
+
1098
+ def get_akpool_rpm(
1099
+ self,
1100
+ *,
1101
+ model: str,
1102
+ ak_id: Optional[str] = None,
1103
+ masked_ak: Optional[str] = None,
1104
+ time_range: Optional[str] = None,
1105
+ from_time: Optional[str] = None,
1106
+ to_time: Optional[str] = None,
1107
+ timeout: TimeoutType = None,
1108
+ ) -> dict:
1109
+ """查询托管 AK 在指定模型上的当前 RPM 用量与上限(GET /apis/v1/model-api-keys/rpm-usage)。
1110
+
1111
+ ``ak_id``(uak-xxx) 与 ``masked_ak``(掩码后 6 位) 二选一,都缺或都给
1112
+ 服务端返回 400;``from_time``/``to_time``(RFC3339) 显式指定时覆盖
1113
+ ``time_range``(Go duration,如 5m/1h,服务端默认 5m)。
1114
+ 响应为嵌套结构:``usage.current_rpm``/``usage.sampled_at``、
1115
+ ``rpm.limit``/``rpm.limit_state``/``rpm.synced_at``、``usage_status``;
1116
+ 查未托管 AK 时 ``ak_id`` 为 null。
1117
+ """
1118
+ validate_ak_id(ak_id)
1119
+ params: dict = {"model": model}
1120
+ if ak_id is not None:
1121
+ params["ak_id"] = ak_id
1122
+ if masked_ak is not None:
1123
+ params["masked_ak"] = masked_ak
1124
+ if time_range is not None:
1125
+ params["time_range"] = time_range
1126
+ if from_time is not None:
1127
+ params["from"] = from_time
1128
+ if to_time is not None:
1129
+ params["to"] = to_time
1130
+ return self._get_central("/model-api-keys/rpm-usage", params=params, timeout=timeout)
1131
+
998
1132
  # ==================== Job operations ====================
999
1133
 
1000
1134
  def create_group(
@@ -1150,6 +1284,7 @@ class APIClient:
1150
1284
  experiment: Optional[str] = None,
1151
1285
  ak_id: Optional[str] = None,
1152
1286
  instance_range: Optional[str] = None,
1287
+ profile: Optional[str] = None,
1153
1288
  ) -> dict:
1154
1289
  """Create a job."""
1155
1290
  body = self.build_create_job_body(
@@ -1180,6 +1315,7 @@ class APIClient:
1180
1315
  credential_type=credential_type,
1181
1316
  account_pool=account_pool,
1182
1317
  eval_config=eval_config,
1318
+ profile=profile,
1183
1319
  profile_id=profile_id,
1184
1320
  profile_version=profile_version,
1185
1321
  instances=instances,
@@ -1260,16 +1396,19 @@ class APIClient:
1260
1396
  experiment: Optional[str] = None,
1261
1397
  ak_id: Optional[str] = None,
1262
1398
  instance_range: Optional[str] = None,
1399
+ profile: Optional[str] = None,
1263
1400
  ) -> dict:
1264
1401
  """Build the /jobs request body for job creation."""
1402
+ if profile is not None and (profile_id is not None or profile_version is not None):
1403
+ raise ValueError("profile cannot be combined with profile_id or profile_version")
1265
1404
  _validate_credential_source_selection(
1266
1405
  account_pool=account_pool,
1267
1406
  credential_type=credential_type,
1268
1407
  params=params,
1269
1408
  params_list=params_list,
1270
1409
  )
1271
- if not template and not profile_id:
1272
- raise ValueError("Either template or profile_id is required")
1410
+ if not template and not profile_id and profile is None:
1411
+ raise ValueError("Either template, profile or profile_id is required")
1273
1412
  validate_ak_selection(ak_id, params, params_list, overrides)
1274
1413
  body: dict = {}
1275
1414
  if template is not None:
@@ -1347,6 +1486,8 @@ class APIClient:
1347
1486
  body["eval_config"] = eval_config
1348
1487
 
1349
1488
  # Profile mode fields
1489
+ if profile is not None:
1490
+ body["profile"] = profile
1350
1491
  if profile_id is not None:
1351
1492
  body["profile_id"] = profile_id
1352
1493
  if profile_version is not None:
@@ -2024,9 +2165,14 @@ class APIClient:
2024
2165
  return self._delete(f"/concurrency-policies/{quote(policy_id)}")
2025
2166
 
2026
2167
  # ==================== Logs and artifacts ====================
2027
- def cancel_group(self, group_id: str) -> dict:
2168
+ def preview_group_cancel(self, group_id: str) -> dict:
2169
+ """Preview the current Group cancel scope and obtain confirmation if required."""
2170
+ return self._post(f"/groups/{quote(group_id)}/cancel-preview", {})
2171
+
2172
+ def cancel_group(self, group_id: str, confirmation_token: Optional[str] = None) -> dict:
2028
2173
  """Cancel all unfinished jobs in a group."""
2029
- return self._post(f"/groups/{quote(group_id)}/cancel", {})
2174
+ body = {"confirmation_token": confirmation_token} if confirmation_token else {}
2175
+ return self._post(f"/groups/{quote(group_id)}/cancel", body)
2030
2176
 
2031
2177
  def get_job_logs(
2032
2178
  self,
@@ -2212,8 +2358,14 @@ class APIClient:
2212
2358
  next_token: Optional[str] = None,
2213
2359
  pagination: Optional[str] = None,
2214
2360
  include_total: bool = True,
2361
+ download_mode: Optional[Literal["oss", "proxy", "both"]] = None,
2215
2362
  ) -> dict:
2216
- """Get one page of artifact download links for a group."""
2363
+ """Get one page of archive links.
2364
+
2365
+ Omit download_mode for the server default and legacy response. Explicit
2366
+ oss/proxy/both modes add a downloads mapping to each artifact; both can
2367
+ return only one successful provider. This method does not download files.
2368
+ """
2217
2369
  params: dict[str, object] = {"skip": skip, "limit": limit}
2218
2370
  if include_post_process:
2219
2371
  params["include_post_process"] = True
@@ -2223,19 +2375,28 @@ class APIClient:
2223
2375
  params["pagination"] = pagination
2224
2376
  if not include_total:
2225
2377
  params["include_total"] = False
2378
+ if download_mode is not None:
2379
+ params["download_mode"] = download_mode
2226
2380
  endpoint = f"/groups/{quote(group_id)}/artifacts"
2227
2381
  response = self._get(endpoint, params=params)
2228
2382
  _ensure_next_token_honored(response, next_token=next_token, endpoint=endpoint)
2229
2383
  return response
2230
2384
 
2231
- def get_group_artifacts(self, group_id: str, include_post_process: bool = False) -> dict:
2232
- """Get artifact download links for a group."""
2385
+ def get_group_artifacts(
2386
+ self,
2387
+ group_id: str,
2388
+ include_post_process: bool = False,
2389
+ *,
2390
+ download_mode: Optional[Literal["oss", "proxy", "both"]] = None,
2391
+ ) -> dict:
2392
+ """Collect all archive-link pages, preserving download_mode on every page."""
2233
2393
  next_token: str | None = None
2234
2394
  skip = 0
2235
2395
  mode: str | None = None
2236
2396
  artifacts: list[dict] = []
2237
2397
  last_page: dict | None = None
2238
2398
 
2399
+ download_options = {"download_mode": download_mode} if download_mode is not None else {}
2239
2400
  while True:
2240
2401
  page = self.get_group_artifacts_page(
2241
2402
  group_id,
@@ -2245,6 +2406,7 @@ class APIClient:
2245
2406
  next_token=next_token,
2246
2407
  pagination="cursor" if mode in {None, "cursor"} else None,
2247
2408
  include_total=False,
2409
+ **download_options,
2248
2410
  )
2249
2411
  last_page = page
2250
2412
  page_artifacts = page.get("artifacts") or []
@@ -2278,9 +2440,17 @@ class APIClient:
2278
2440
  )
2279
2441
  return result
2280
2442
 
2281
- def get_job_artifacts(self, job_ids: list) -> list:
2282
- """Get artifact download links for one or more jobs."""
2283
- return self._post("/jobs/artifacts", {"job_ids": job_ids})
2443
+ def get_job_artifacts(
2444
+ self,
2445
+ job_ids: list,
2446
+ *,
2447
+ download_mode: Optional[Literal["oss", "proxy", "both"]] = None,
2448
+ ) -> list:
2449
+ """Get archive links; explicit modes also return typed ``downloads`` entries."""
2450
+ body = {"job_ids": job_ids}
2451
+ if download_mode is not None:
2452
+ body["download_mode"] = download_mode
2453
+ return self._post("/jobs/artifacts", body)
2284
2454
 
2285
2455
  def get_job_artifacts_manifest(
2286
2456
  self,
@@ -2295,6 +2465,52 @@ class APIClient:
2295
2465
  params=params,
2296
2466
  )
2297
2467
 
2468
+ def wait_for_job_artifacts_manifest(self, job_id: str) -> dict:
2469
+ """Poll the new API within one per-Job budget; HTTP failures keep normal retries."""
2470
+ budget = self.config.artifact_manifest_wait_timeout
2471
+ if not math.isfinite(budget) or budget <= 0:
2472
+ raise ConfigurationError(
2473
+ "AP_ARTIFACT_MANIFEST_WAIT_TIMEOUT must be a finite positive number"
2474
+ )
2475
+ deadline = time.monotonic() + budget
2476
+ while True:
2477
+ response, request_id = self._request(
2478
+ "GET",
2479
+ f"/jobs/{quote(job_id, safe='')}/artifacts/manifest",
2480
+ deadline=deadline,
2481
+ _return_response=True,
2482
+ )
2483
+ try:
2484
+ payload = response.json()
2485
+ except ValueError as exc:
2486
+ raise APIError(
2487
+ response.status_code, "Invalid artifact manifest response", request_id, response
2488
+ ) from exc
2489
+ if not isinstance(payload, dict) or payload.get("job_id") != job_id:
2490
+ raise APIError(
2491
+ response.status_code, "Invalid artifact manifest response", request_id, response
2492
+ )
2493
+ if response.status_code == 200:
2494
+ if "manifest" not in payload or (
2495
+ payload["manifest"] is not None and not isinstance(payload["manifest"], list)
2496
+ ):
2497
+ raise APIError(200, "Invalid artifact manifest response", request_id, response)
2498
+ return payload
2499
+ if response.status_code != 202 or payload.get("status") != "preparing":
2500
+ raise APIError(
2501
+ response.status_code,
2502
+ "Invalid artifact preparation response",
2503
+ request_id,
2504
+ response,
2505
+ )
2506
+ try:
2507
+ delay = float(response.headers.get("Retry-After", "2"))
2508
+ except ValueError:
2509
+ delay = 2.0
2510
+ if not math.isfinite(delay) or delay <= 0:
2511
+ delay = 2.0
2512
+ _deadline_sleep(delay, deadline)
2513
+
2298
2514
  def get_artifact_file_download_urls(self, job_id: str, paths: list[str]) -> dict:
2299
2515
  """Get download URLs for multiple artifact files."""
2300
2516
  endpoint = f"/jobs/{quote(job_id, safe='')}/artifacts/files/download"
@@ -3208,3 +3424,34 @@ def _summarize_json_value(data: Any, limit: int) -> str:
3208
3424
  summary = {"_count": len(data), "_sample": sample}
3209
3425
  return _truncate_text(_safe_json(summary), limit)
3210
3426
  return _truncate_text(_safe_json(data), limit)
3427
+
3428
+
3429
+ def _manifest_budget_remaining(deadline: float) -> float:
3430
+ remaining = deadline - time.monotonic()
3431
+ if remaining <= 0:
3432
+ raise APIError(None, "Artifact manifest wait timed out", None, None)
3433
+ return remaining
3434
+
3435
+
3436
+ def _deadline_timeout(
3437
+ timeout: TimeoutType, deadline: Optional[float]
3438
+ ) -> Union[TimeoutType, HTTPTimeout]:
3439
+ if deadline is None:
3440
+ return timeout
3441
+ remaining = _manifest_budget_remaining(deadline)
3442
+
3443
+ def bounded(value):
3444
+ return min(value, remaining) if value is not None else remaining
3445
+
3446
+ if isinstance(timeout, tuple):
3447
+ connect, read = (bounded(value) for value in timeout)
3448
+ else:
3449
+ connect = read = bounded(timeout)
3450
+ # The connect phase must consume the same remaining budget as the read.
3451
+ return HTTPTimeout(total=remaining, connect=connect, read=read)
3452
+
3453
+
3454
+ def _deadline_sleep(delay: float, deadline: Optional[float]) -> None:
3455
+ if deadline is not None:
3456
+ delay = min(delay, _manifest_budget_remaining(deadline))
3457
+ time.sleep(delay)