feat(credential): 暴露 use_sts_credentials / use_sts_from_headers 供非 server 场景手动注入 STS

StsRefreshMiddleware 只在 agentrun server 内生效；用户用自有 Web 框架（FastAPI/
Flask/Django）或非 HTTP 任务时无法被拦截。新增两个公开上下文管理器（从 agentrun
顶层导出），让用户手动注入最新 STS：
- use_sts_credentials(ak, sk, sts)：显式传值
- use_sts_from_headers(headers)：从任意请求头映射解析（同 x-fc-* 约定，大小写不敏感）

二者写入同一请求级 overlay，刷新逻辑与中间件完全一致（含三元组齐全校验、自动复位）。
中间件本身重构为直接复用 use_sts_from_headers —— 单一真相，server 内外同一套逻辑。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: OhYee <oyohyee@oyohyee.com>

Author: OhYee

AGENTS.md

@@ -481,7 +481,19 @@ STS 临时凭证（ak/sk/security_token）会过期。部署在函数计算（FC
 1. **请求级 overlay**：`agentrun/server/sts_middleware.py` 解析 FC 头
    （默认 `x-fc-access-key-id` / `x-fc-access-key-secret` / `x-fc-security-token`，
    可经构造参数或 `AGENTRUN_STS_HEADER_*` 环境变量覆盖），写入
-   `agentrun/utils/credential_context.py` 的 `contextvars` overlay。
+   `agentrun/utils/credential_context.py` 的 `contextvars` overlay。中间件本身
+   只是 `use_sts_from_headers` 的薄封装（加 FC 门控），二者共用同一套解析逻辑。
+
+   **非 agentrun server 场景**（自有 FastAPI / Flask / Django、或非 HTTP 任务）：
+   中间件不会运行，需用户手动注入。SDK 顶层导出两个上下文管理器：
+   - `agentrun.use_sts_credentials(ak, sk, sts)` —— 显式传值；
+   - `agentrun.use_sts_from_headers(headers)` —— 从任意请求头映射解析（同 `x-fc-*`）。
+
+   ```python
+   from agentrun import use_sts_from_headers
+   with use_sts_from_headers(request.headers):
+       ...  # 块内所有 SDK 调用使用最新 STS，退出自动复位
+   ```
 
 2. **Config 懒解析**：`Config` 的三个凭证 getter 按
    **显式传入 > 请求级 overlay（仅当三者均未显式传入）> 环境变量** 解析。

agentrun/__init__.py

@@ -135,6 +135,11 @@
 # ToolSet
 from agentrun.toolset import ToolSet, ToolSetClient
 from agentrun.utils.config import Config
+from agentrun.utils.credential_context import (
+    StsCredential,
+    use_sts_credentials,
+    use_sts_from_headers,
+)
 from agentrun.utils.exception import (
     ResourceAlreadyExistError,
     ResourceNotExistError,
@@ -335,6 +340,10 @@
     "ResourceNotExistError",
     "ResourceAlreadyExistError",
     "Config",
+    ######## STS 凭证刷新（非 server 场景手动注入） ########
+    "StsCredential",
+    "use_sts_credentials",
+    "use_sts_from_headers",
 ]
 
 # Memory Collection 模块的所有导出（延迟加载）

agentrun/server/sts_middleware.py

@@ -33,16 +33,7 @@
 from starlette.responses import Response
 from starlette.types import ASGIApp
 
-from agentrun.utils.credential_context import (
-    StsCredential,
-    reset_request_sts,
-    set_request_sts,
-)
-from agentrun.utils.log import logger
-
-DEFAULT_ACCESS_KEY_ID_HEADER = "x-fc-access-key-id"
-DEFAULT_ACCESS_KEY_SECRET_HEADER = "x-fc-access-key-secret"
-DEFAULT_SECURITY_TOKEN_HEADER = "x-fc-security-token"
+from agentrun.utils.credential_context import use_sts_from_headers
 
 
 def _detect_enabled() -> bool:
@@ -69,61 +60,21 @@ def __init__(
         super().__init__(app)
         # enabled=None 时自动探测（FC 环境或显式环境变量开关）。
         self._enabled = _detect_enabled() if enabled is None else enabled
-        self._ak_header = (
-            access_key_id_header
-            or os.getenv("AGENTRUN_STS_HEADER_ACCESS_KEY_ID")
-            or DEFAULT_ACCESS_KEY_ID_HEADER
-        )
-        self._sk_header = (
-            access_key_secret_header
-            or os.getenv("AGENTRUN_STS_HEADER_ACCESS_KEY_SECRET")
-            or DEFAULT_ACCESS_KEY_SECRET_HEADER
-        )
-        self._sts_header = (
-            security_token_header
-            or os.getenv("AGENTRUN_STS_HEADER_SECURITY_TOKEN")
-            or DEFAULT_SECURITY_TOKEN_HEADER
-        )
+        # 头名解析（参数 > 环境变量 > 默认）交由 sts_from_headers 处理，这里只存原值。
+        self._ak_header = access_key_id_header
+        self._sk_header = access_key_secret_header
+        self._sts_header = security_token_header
 
     async def dispatch(self, request: Request, call_next) -> Response:
         if not self._enabled:
             return await call_next(request)
 
-        cred = self._extract(request)
-        if cred is None:
-            return await call_next(request)
-
-        token = set_request_sts(cred)
-        try:
+        # 直接复用公开上下文管理器：解析请求头 -> 注入 overlay -> 退出复位。
+        # 三元组不齐全时 use_sts_from_headers 不覆盖（透传），与手动注入完全一致。
+        with use_sts_from_headers(
+            request.headers,
+            access_key_id_header=self._ak_header,
+            access_key_secret_header=self._sk_header,
+            security_token_header=self._sts_header,
+        ):
             return await call_next(request)
-        finally:
-            reset_request_sts(token)
-
-    def _extract(self, request: Request) -> Optional[StsCredential]:
-        # starlette Headers 大小写不敏感。
-        headers = request.headers
-        cred = StsCredential(
-            access_key_id=headers.get(self._ak_header),
-            access_key_secret=headers.get(self._sk_header),
-            security_token=headers.get(self._sts_header),
-        )
-
-        # 只有三元组齐全才视为有效 STS 刷新，避免把新 sts 与陈旧 ak/sk 混用。
-        if not cred.is_complete():
-            if not (
-                cred.access_key_id
-                or cred.access_key_secret
-                or cred.security_token
-            ):
-                return None  # 无任何凭证头：常规非 FC 请求，静默跳过。
-            logger.warning(
-                "STS headers incomplete (ak=%s, sk=%s, sts=%s); ignoring"
-                " partial credential set and falling back to env.",
-                "set" if cred.access_key_id else "unset",
-                "set" if cred.access_key_secret else "unset",
-                "set" if cred.security_token else "unset",
-            )
-            return None
-
-        logger.debug("STS overlay applied from request headers")
-        return cred

agentrun/utils/credential_context.py

@@ -22,9 +22,16 @@
 
 from __future__ import annotations
 
+import contextlib
 import contextvars
+import os
 from dataclasses import dataclass
-from typing import Optional
+from typing import Iterator, Mapping, Optional
+
+# FC 注入 STS 的默认头名（可经构造参数或环境变量覆盖）。
+DEFAULT_ACCESS_KEY_ID_HEADER = "x-fc-access-key-id"
+DEFAULT_ACCESS_KEY_SECRET_HEADER = "x-fc-access-key-secret"
+DEFAULT_SECURITY_TOKEN_HEADER = "x-fc-security-token"
 
 
 @dataclass(frozen=True)
@@ -78,3 +85,109 @@ def reset_request_sts(token: contextvars.Token) -> None:
 def get_request_sts() -> Optional[StsCredential]:
     """获取当前请求的 STS 覆盖层；未设置时返回 ``None``。"""
     return _current_sts.get()
+
+
+def _resolve_header_name(
+    explicit: Optional[str], env_key: str, default: str
+) -> str:
+    """解析头名：构造参数 > 环境变量 > 默认值；统一转小写。"""
+    return (explicit or os.getenv(env_key) or default).lower()
+
+
+def sts_from_headers(
+    headers: Mapping[str, str],
+    *,
+    access_key_id_header: Optional[str] = None,
+    access_key_secret_header: Optional[str] = None,
+    security_token_header: Optional[str] = None,
+) -> Optional[StsCredential]:
+    """从请求头映射解析 STS 三元组；不齐全则返回 ``None``。
+
+    仅当 ak/sk/sts 三者齐全才视为有效刷新（避免把新 sts 与陈旧/环境变量里的
+    ak/sk 混用）。``headers`` 可为任意 Mapping（如 ``dict`` 或 Starlette
+    ``Headers``），按头名**大小写不敏感**查找。头名优先级：参数 > 环境变量
+    （``AGENTRUN_STS_HEADER_*``）> 默认（``x-fc-*``）。
+    """
+    ak_name = _resolve_header_name(
+        access_key_id_header,
+        "AGENTRUN_STS_HEADER_ACCESS_KEY_ID",
+        DEFAULT_ACCESS_KEY_ID_HEADER,
+    )
+    sk_name = _resolve_header_name(
+        access_key_secret_header,
+        "AGENTRUN_STS_HEADER_ACCESS_KEY_SECRET",
+        DEFAULT_ACCESS_KEY_SECRET_HEADER,
+    )
+    sts_name = _resolve_header_name(
+        security_token_header,
+        "AGENTRUN_STS_HEADER_SECURITY_TOKEN",
+        DEFAULT_SECURITY_TOKEN_HEADER,
+    )
+
+    lower = {str(k).lower(): v for k, v in headers.items()}
+    cred = StsCredential(
+        access_key_id=lower.get(ak_name),
+        access_key_secret=lower.get(sk_name),
+        security_token=lower.get(sts_name),
+    )
+    return cred if cred.is_complete() else None
+
+
+@contextlib.contextmanager
+def use_sts_credentials(
+    access_key_id: Optional[str] = None,
+    access_key_secret: Optional[str] = None,
+    security_token: Optional[str] = None,
+) -> Iterator[StsCredential]:
+    """在 ``with`` 块内临时使用给定 STS 临时凭证（请求级 overlay），退出自动复位。
+
+    适用于**不经过 agentrun server** 的场景：自有 FastAPI / Flask / Django，或
+    非 HTTP 的任务里，从上游 / 请求头拿到最新 STS 后注入——块内所有 SDK 调用
+    （以及其内创建的 asyncio 任务）即使用这组凭证。
+
+    Examples:
+        >>> with use_sts_credentials(ak, sk, sts):
+        ...     knowledgebase.retrieve(...)  # 使用最新 STS
+
+    Note:
+        基于 ``contextvars``，按当前任务/线程隔离；用户自行 ``threading.Thread``
+        起的裸线程不会继承（``asyncio.create_task`` 会）。
+    """
+    cred = StsCredential(access_key_id, access_key_secret, security_token)
+    token = set_request_sts(cred)
+    try:
+        yield cred
+    finally:
+        reset_request_sts(token)
+
+
+@contextlib.contextmanager
+def use_sts_from_headers(
+    headers: Mapping[str, str],
+    *,
+    access_key_id_header: Optional[str] = None,
+    access_key_secret_header: Optional[str] = None,
+    security_token_header: Optional[str] = None,
+) -> Iterator[Optional[StsCredential]]:
+    """从请求头映射解析 STS 并在 ``with`` 块内生效；三元组不齐全则不覆盖（透传）。
+
+    与 :class:`agentrun.server.sts_middleware.StsRefreshMiddleware` 共用同一套
+    解析逻辑。适用于在自有 Web 框架里手动接入：
+
+        >>> with use_sts_from_headers(request.headers):
+        ...     await invoke_agent(request)
+    """
+    cred = sts_from_headers(
+        headers,
+        access_key_id_header=access_key_id_header,
+        access_key_secret_header=access_key_secret_header,
+        security_token_header=security_token_header,
+    )
+    if cred is None:
+        yield None
+        return
+    token = set_request_sts(cred)
+    try:
+        yield cred
+    finally:
+        reset_request_sts(token)

tests/unittests/test_sts_refresh.py

@@ -18,6 +18,8 @@
     get_request_sts,
     reset_request_sts,
     set_request_sts,
+    use_sts_credentials,
+    use_sts_from_headers,
 )
 
 
@@ -418,3 +420,73 @@ async def run():
         if e.event == EventType.TEXT
     )
     assert text == "ak=OV_AK;sts=OV_STS", text
+
+
+# --------------------------------------------------------------------------- #
+# 公开 API：非 server 场景手动注入
+# --------------------------------------------------------------------------- #
+def test_use_sts_credentials_context_manager(monkeypatch):
+    monkeypatch.setenv("AGENTRUN_ACCESS_KEY_ID", "ENV_AK")
+    cfg = Config()
+    assert cfg.get_access_key_id() == "ENV_AK"
+    with use_sts_credentials("OV_AK", "OV_SK", "OV_STS"):
+        assert cfg.get_access_key_id() == "OV_AK"
+        assert cfg.get_access_key_secret() == "OV_SK"
+        assert cfg.get_security_token() == "OV_STS"
+    # 退出自动复位
+    assert cfg.get_access_key_id() == "ENV_AK"
+    assert get_request_sts() is None
+
+
+def test_use_sts_from_headers_complete_case_insensitive():
+    cfg = Config()
+    headers = {
+        "X-Fc-Access-Key-Id": "H_AK",  # 大小写不敏感
+        "x-fc-access-key-secret": "H_SK",
+        "X-FC-SECURITY-TOKEN": "H_STS",
+    }
+    with use_sts_from_headers(headers) as cred:
+        assert cred is not None
+        assert cfg.get_access_key_id() == "H_AK"
+        assert cfg.get_security_token() == "H_STS"
+    assert get_request_sts() is None
+
+
+def test_use_sts_from_headers_partial_no_override(monkeypatch):
+    monkeypatch.setenv("AGENTRUN_ACCESS_KEY_ID", "ENV_AK")
+    cfg = Config()
+    # 只有 sts、缺 ak/sk -> 不构成完整三元组 -> 不覆盖
+    with use_sts_from_headers({"x-fc-security-token": "H_STS"}) as cred:
+        assert cred is None
+        assert cfg.get_access_key_id() == "ENV_AK"
+        assert cfg.get_security_token() == ""
+    assert get_request_sts() is None
+
+
+def test_sts_from_headers_helper():
+    from agentrun.utils.credential_context import sts_from_headers
+
+    assert sts_from_headers({"x-fc-access-key-id": "a"}) is None  # 不齐全
+    cred = sts_from_headers({
+        "x-fc-access-key-id": "a",
+        "x-fc-access-key-secret": "b",
+        "x-fc-security-token": "c",
+    })
+    assert cred is not None
+    assert (
+        cred.access_key_id,
+        cred.access_key_secret,
+        cred.security_token,
+    ) == ("a", "b", "c")
+
+
+def test_public_exports_available():
+    import agentrun
+
+    for name in (
+        "StsCredential",
+        "use_sts_credentials",
+        "use_sts_from_headers",
+    ):
+        assert hasattr(agentrun, name), f"{name} not exported"
+        assert name in agentrun.__all__, f"{name} missing from __all__"