agentengine-sdk-python 0.6.9__py3-none-any.whl → 0.8.0__py3-none-any.whl

ksadk/model_proxy/namespace.py

@@ -0,0 +1,173 @@
+"""codex namespace 工具拍平 + 还原(responses↔chat 路径)。
+
+codex 0.142+ 用 namespace 工具声明插件/MCP 工具:
+
+    {"type":"namespace","name":"mcp__x__","tools":[{"type":"function","name":"foo",...}]}
+
+并在 input 历史里发 namespace-qualified function_call:
+
+    {"type":"function_call","name":"foo","namespace":"mcp__x__",...}
+
+chat completions 只认顶层 ``function`` 工具(chat 嵌套)和 ``function.name``,
+不认 ``namespace`` 字段。本模块把 namespace 工具**拍平**成顶层 function,
+名字用确定性的 ``<namespace>__<child>``(超长 sha256 截断,与 cc-switch 一致),
+响应侧再把 flat 名字**还原**回 ``{name, namespace}``,让 codex client 能匹配
+自己的 namespace 工具注册表。
+
+请求侧拍平 + 响应侧还原都从**同一份 request tools** 经
+:func:`flatten_namespace_tool_name` 派生名字映射,前后一致,无需跨请求状态。
+
+参照 cc-switch ``transform_codex_responses_namespace.rs``(用 Python 重写)。
+"""
+
+from __future__ import annotations
+
+import hashlib
+from typing import Any
+
+CHAT_TOOL_NAME_MAX_LEN = 64
+
+
+def flatten_namespace_tool_name(namespace: str, name: str) -> str:
+    """``<namespace>__<name>``;超长则尾部接 sha256 截断,保持在 64 字符内。
+
+    截断按字符边界,保证返回值是合法 UTF-8(不会从多字节字符中间切断)。
+    """
+    full = f"{namespace}__{name}"
+    if len(full) <= CHAT_TOOL_NAME_MAX_LEN:
+        return full
+    suffix = "__" + hashlib.sha256(full.encode("utf-8")).hexdigest()[:8]
+    prefix_len = CHAT_TOOL_NAME_MAX_LEN - len(suffix)
+    prefix = ""
+    for ch in full:
+        if len(prefix.encode("utf-8")) + len(ch.encode("utf-8")) > prefix_len:
+            break
+        prefix += ch
+    return prefix + suffix
+
+
+def _namespace_children(tool: dict) -> list[dict]:
+    for key in ("tools", "children"):
+        kids = tool.get(key)
+        if isinstance(kids, list):
+            return kids
+    return []
+
+
+def build_restore_map(tools: list[Any] | None) -> dict[str, dict[str, str]]:
+    """从 request 的 namespace 工具声明建 ``flat -> {namespace, name}`` 映射。
+
+    两个不同 child 拍平后撞名时,后者覆盖前者(与 cc-switch ``or_insert`` 不同,
+    cc-switch 用 first-wins 并对撞名报错;这里取 first-wins + 调用方撞名检测)。
+    """
+    restore: dict[str, dict[str, str]] = {}
+    if not tools:
+        return restore
+    for tool in tools:
+        if not isinstance(tool, dict) or tool.get("type") != "namespace":
+            continue
+        namespace = (tool.get("name") or "").strip()
+        if not namespace:
+            continue
+        for child in _namespace_children(tool):
+            if not isinstance(child, dict) or child.get("type") != "function":
+                continue
+            name = (child.get("name") or "").strip()
+            if not name:
+                continue
+            flat = flatten_namespace_tool_name(namespace, name)
+            restore.setdefault(flat, {"namespace": namespace, "name": name})
+    return restore
+
+
+def _rewrite_qualified_calls(value: Any, owners: dict[str, dict[str, str]]) -> bool:
+    """递归把 input 里的 namespace-qualified function_call 重写成 flat name。
+
+    ``{type:function_call, name:child, namespace:ns}`` -> ``{type:function_call, name:flat}``
+    (去掉 namespace 字段);只在 (ns, name) 与 owners 条目匹配时重写。
+    """
+    changed = False
+    if isinstance(value, list):
+        for item in value:
+            changed |= _rewrite_qualified_calls(item, owners)
+    elif isinstance(value, dict):
+        if value.get("type") == "function_call":
+            namespace = (value.get("namespace") or "").strip()
+            name = (value.get("name") or "").strip()
+            if namespace and name:
+                flat = flatten_namespace_tool_name(namespace, name)
+                entry = owners.get(flat)
+                if entry and entry["namespace"] == namespace and entry["name"] == name:
+                    value["name"] = flat
+                    value.pop("namespace", None)
+                    changed = True
+        for child in value.values():
+            changed |= _rewrite_qualified_calls(child, owners)
+    return changed
+
+
+def flatten_request_namespaces(body: dict) -> dict[str, dict[str, str]]:
+    """原地拍平 body 的 namespace 工具,返回 restore_map(供响应侧还原)。
+
+    - tools:namespace children 提升为顶层 responses function 工具(flat name),
+      原 namespace 工具移除;检测撞名(两个 child 拍平后同名)则 raise。
+    - input:namespace-qualified function_call 重写为 flat name。
+    - tool_choice:namespace-typed 丢弃(chat 不认)。
+
+    返回 ``flat -> {namespace, name}`` 映射;无 namespace 工具时返回空 dict,
+    调用方可据此跳过响应侧还原。
+    """
+    tools = body.get("tools")
+    if not isinstance(tools, list):
+        return {}
+    restore = build_restore_map(tools)
+    flat_tools: list[dict] = []
+    seen_flat: set[str] = set()
+    for tool in tools:
+        if not isinstance(tool, dict):
+            continue
+        if tool.get("type") == "namespace":
+            namespace = (tool.get("name") or "").strip()
+            for child in _namespace_children(tool):
+                if not isinstance(child, dict) or child.get("type") != "function":
+                    continue
+                name = (child.get("name") or "").strip()
+                if not name or not namespace:
+                    continue
+                flat = flatten_namespace_tool_name(namespace, name)
+                if flat in seen_flat:
+                    raise ValueError(
+                        f"namespace 拍平撞名:{flat} 来自不同 child,上游无法消歧"
+                    )
+                seen_flat.add(flat)
+                flat_tools.append({
+                    "type": "function",
+                    "name": flat,
+                    "description": child.get("description", ""),
+                    "parameters": child.get("parameters", {"type": "object", "properties": {}}),
+                    **({"strict": child["strict"]} if "strict" in child else {}),
+                })
+        else:
+            flat_tools.append(tool)
+    body["tools"] = flat_tools
+    if restore and isinstance(body.get("input"), list):
+        _rewrite_qualified_calls(body["input"], restore)
+    tc = body.get("tool_choice")
+    if isinstance(tc, dict) and tc.get("type") == "namespace":
+        body.pop("tool_choice", None)
+    return restore
+
+
+def restore_function_call(item: dict, restore_map: dict[str, dict[str, str]]) -> dict:
+    """把响应里 function_call 的 flat name 还原成 ``{name, namespace}``(原地改)。
+
+    无映射或 name 不在映射中时原样返回(普通 function 工具的 call 不受影响)。
+    """
+    if not restore_map or item.get("type") != "function_call":
+        return item
+    flat = item.get("name")
+    entry = restore_map.get(flat) if isinstance(flat, str) else None
+    if entry:
+        item["name"] = entry["name"]
+        item["namespace"] = entry["namespace"]
+    return item

ksadk/model_proxy/protocol_proxy.py

@@ -0,0 +1,22 @@
+"""可独立运行的转换 proxy 入口:读 env → ProxyConfig → uvicorn。
+
+用法:
+  KSPMAS_API_KEY=xxx python -m ksadk.model_proxy.protocol_proxy   # 默认 127.0.0.1:8899
+"""
+
+import os
+
+import uvicorn
+
+from .config import ProxyConfig
+from .server import create_app
+
+
+def main():
+    config = ProxyConfig.from_env()
+    port = int(os.environ.get("PORT", "8899"))
+    uvicorn.run(create_app(config), host="127.0.0.1", port=port, log_level="warning")
+
+
+if __name__ == "__main__":
+    main()

ksadk/model_proxy/server.py

@@ -0,0 +1,300 @@
+"""app factory + 进程内转换 proxy 生命周期管理。
+
+把 HTTP 层与配置/转换解耦:create_app(config) 返回 FastAPI app;
+ProxyServer 在后台线程懒起一个 localhost uvicorn,幂等 + 加锁 + 干净回收(随 runner 退出)。
+"""
+
+import asyncio
+import json
+import logging
+import socket
+import threading
+import time
+import uuid
+
+import httpx
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from .config import _LOOPBACK_HOSTS, ProxyConfig
+from .transform import Streamer, UnsupportedToolsError, chat_to_response, responses_to_chat
+
+logger = logging.getLogger(__name__)
+
+
+def _upstream_error(status_code: int) -> JSONResponse:
+    """Return a stable public error without relaying an upstream traceback."""
+    return JSONResponse(
+        status_code=status_code,
+        content={
+            "error": {
+                "type": "upstream_error",
+                "message": "The model upstream rejected the request.",
+            }
+        },
+    )
+
+
+def create_app(config: ProxyConfig) -> FastAPI:
+    base = config.upstream_base.rstrip("/")
+    headers = {"Authorization": f"Bearer {config.api_key}", "Content-Type": "application/json"}
+    app = FastAPI()
+    # 应用层取消信号:stop() 时置位,让活动 SSE 的 _stream_gen 主动 break(比单靠 force_exit 可靠)
+    cancel = threading.Event()
+    app.state.cancel_event = cancel
+
+    @app.get("/healthz")
+    async def healthz():
+        return {"ok": True, "upstream": base}
+
+    @app.post("/v1/chat/completions")
+    async def chat_passthrough(req: Request):
+        """chat 直通端点(env 框架发 chat 经代理原样转发上游,不转换)。
+
+        让代理成为统一模型出口:env 框架(ADK/langchain/deepagents 发 chat)经此直通,
+        codex(发 responses)经 /v1/responses 转换。统一出口为后续能力探测/路由打底。
+        chat 直通字节级转发 + 流式透传,前缀缓存不受影响。
+        """
+        if (
+            config.local_token
+            and req.headers.get("authorization", "") != f"Bearer {config.local_token}"
+        ):
+            return JSONResponse(
+                status_code=401, content={"error": {"type": "authentication_error"}}
+            )
+        body = await req.json()
+        is_stream = bool(body.get("stream"))
+        up_headers = {
+            "Authorization": f"Bearer {config.api_key}",
+            "Content-Type": "application/json",
+        }
+        if is_stream:
+            return StreamingResponse(
+                _chat_passthrough_stream(config, base, up_headers, body),
+                media_type="text/event-stream",
+                headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+            )
+        async with httpx.AsyncClient(timeout=config.timeout) as c:
+            r = await c.post(f"{base}/chat/completions", json=body, headers=up_headers)
+            if r.is_error:
+                logger.warning("chat upstream rejected request: status=%s", r.status_code)
+                return _upstream_error(r.status_code)
+            return JSONResponse(status_code=r.status_code, content=r.json() if r.content else None)
+
+    @app.post("/v1/responses")
+    async def responses(req: Request):
+        # 本地鉴权:codex -> 本地 proxy 用 KSADK_PROXY_TOKEN;上游凭证独立保存于 config.api_key
+        if (
+            config.local_token
+            and req.headers.get("authorization", "") != f"Bearer {config.local_token}"
+        ):
+            return JSONResponse(
+                status_code=401,
+                content={
+                    "error": {"type": "authentication_error", "message": "invalid proxy token"}
+                },
+            )
+        body = await req.json()
+        try:
+            chat_req, restore_map = responses_to_chat(body)
+        except UnsupportedToolsError as e:
+            logger.info("responses request uses unsupported tools: %s", e)
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": {
+                        "type": "unsupported_tools",
+                        "message": "The request uses tools unsupported by the model upstream.",
+                    }
+                },
+            )
+        rid = "resp_" + uuid.uuid4().hex[:24]
+        model = body.get("model")
+        if body.get("stream"):
+            chat_req["stream"] = True
+            chat_req["stream_options"] = {"include_usage": True}
+            return StreamingResponse(
+                _stream_gen(config, base, headers, chat_req, rid, model, cancel, restore_map),
+                media_type="text/event-stream",
+                headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+            )
+        async with httpx.AsyncClient(timeout=config.timeout) as c:
+            r = await c.post(f"{base}/chat/completions", json=chat_req, headers=headers)
+            if r.status_code != 200:
+                logger.warning("responses upstream rejected request: status=%s", r.status_code)
+                return _upstream_error(r.status_code)
+            return JSONResponse(chat_to_response(r.json(), rid, restore_map))
+
+    return app
+
+
+async def _chat_passthrough_stream(config: ProxyConfig, base: str, headers: dict, body: dict):
+    """chat SSE 字节级透传(不转换),让前缀缓存/流式格式与直连一致。"""
+    try:
+        async with httpx.AsyncClient(timeout=config.timeout) as c:
+            async with c.stream(
+                "POST", f"{base}/chat/completions", json=body, headers=headers
+            ) as r:
+                async for line in r.aiter_lines():
+                    yield line + "\n"
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("chat upstream stream failed", exc_info=exc)
+        yield 'data: {"error":"The model upstream stream failed."}\n\n'
+
+
+async def _stream_gen(
+    config: ProxyConfig,
+    base: str,
+    headers: dict,
+    chat_req: dict,
+    rid: str,
+    model,
+    cancel: threading.Event,
+    restore_map: dict | None = None,
+):
+    s = Streamer(rid, model, restore_map)
+    for e in s.start():
+        yield e
+    try:
+        # 超时必须有界(config.timeout):无限超时的活动 SSE 会让 stop() 后线程仍存活
+        async with httpx.AsyncClient(timeout=config.timeout) as c:
+            async with c.stream(
+                "POST", f"{base}/chat/completions", json=chat_req, headers=headers
+            ) as r:
+                if r.status_code != 200:
+                    logger.warning("responses upstream stream rejected: status=%s", r.status_code)
+                    yield Streamer.ev(
+                        "response.failed",
+                        {
+                            "type": "response.failed",
+                            "response": {
+                                **s._resp("failed"),
+                                "error": {"message": "The model upstream rejected the request."},
+                            },
+                        },
+                    )
+                    return
+                async for line in r.aiter_lines():
+                    if cancel.is_set():
+                        return  # stop() 置位:主动中断活动 SSE,让线程能及时回收
+                    line = line.strip()
+                    if not line.startswith("data:"):
+                        continue
+                    data = line[5:].strip()
+                    if data == "[DONE]":
+                        break
+                    try:
+                        chunk = json.loads(data)
+                    except ValueError:
+                        continue
+                    for e in s.handle(chunk):
+                        yield e
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("responses upstream stream failed", exc_info=exc)
+        yield Streamer.ev(
+            "response.failed",
+            {
+                "type": "response.failed",
+                "response": {
+                    **s._resp("failed"),
+                    "error": {"message": "The model upstream stream failed."},
+                },
+            },
+        )
+        return
+    for e in s.finalize():
+        yield e
+
+
+class ProxyServer:
+    """进程内 localhost 转换 proxy 的生命周期管理。
+
+    - start 幂等(已启动直接返回)+ 加锁,避免重复 start 覆盖 server/thread 造成泄漏;
+    - 端口在 start 时 bind 并**一直持有 socket**(传给 uvicorn),避免"先释放再重绑"的 TOCTOU;
+    - 启动超时与 stop 都会干净回收线程与 socket。
+    """
+
+    def __init__(self, config: ProxyConfig, host: str = "127.0.0.1", port: int = 0):
+        if host not in _LOOPBACK_HOSTS and not config.local_token:
+            raise ValueError(
+                f"非回环监听(host={host})必须设置 local_token,否则是无鉴权代持上游凭证的开放代理"
+            )
+        self.config = config
+        self.host = host
+        self._requested_port = port
+        self.port = port
+        self._lock = threading.Lock()
+        self._sock: socket.socket | None = None
+        self._server: uvicorn.Server | None = None
+        self._thread: threading.Thread | None = None
+        self._app: FastAPI | None = None
+
+    @property
+    def base_url(self) -> str:
+        return f"http://{self.host}:{self.port}/v1"
+
+    def start(self, wait_ready: float = 10.0) -> str:
+        with self._lock:
+            if self._server is not None:
+                return self.base_url  # 幂等
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            try:
+                sock.bind((self.host, self._requested_port))
+            except OSError:
+                sock.close()
+                raise
+            self.port = int(sock.getsockname()[1])
+            app = create_app(self.config)
+            self._app = app
+            # timeout_graceful_shutdown 设上限:stop 时挂在 yield 的 SSE generator 不再消费,
+            # 若无限等待连接完成(None 默认)会让 shutdown 卡死、线程回收不掉
+            cfg = uvicorn.Config(app, log_level="warning", timeout_graceful_shutdown=1)
+            self._sock = sock
+            server = uvicorn.Server(cfg)
+            self._server = server
+            # serve(sockets=[sock]) 跨平台:uvicorn 直接复用已绑定 socket,
+            # 不走 Config(fd=...) 那条仅 Unix 可用的 socket.fromfd 路径
+            self._thread = threading.Thread(
+                target=lambda: asyncio.run(server.serve(sockets=[sock])), daemon=True
+            )
+            self._thread.start()
+            deadline = time.time() + wait_ready
+            while time.time() < deadline:
+                try:
+                    if (
+                        httpx.get(f"http://{self.host}:{self.port}/healthz", timeout=1).status_code
+                        == 200
+                    ):
+                        return self.base_url
+                except Exception:  # noqa: BLE001
+                    time.sleep(0.05)
+            self._teardown()  # 启动超时也要清理线程与 socket
+            raise RuntimeError(f"ProxyServer 未在 {wait_ready}s 内就绪")
+
+    def stop(self):
+        with self._lock:
+            self._teardown()
+
+    def _teardown(self):
+        if self._app is not None:
+            # 应用层取消信号:让活动 SSE 的 _stream_gen 主动 break
+            self._app.state.cancel_event.set()
+        if self._server is not None:
+            # should_exit 触发主循环进入 shutdown;force_exit 让 shutdown 跳过等待
+            # connections/tasks 完成(活动 SSE 下不再无限等)。两者缺一:只 should_exit
+            # 会被活动连接卡死,只 force_exit 主循环根本不进 shutdown。
+            self._server.should_exit = True
+            self._server.force_exit = True
+        if self._thread is not None:
+            self._thread.join(timeout=5)
+        if self._sock is not None:
+            try:
+                self._sock.close()
+            except OSError:
+                pass
+        self._server = None
+        self._thread = None
+        self._sock = None
+        self._app = None