union-app-chat-stream 1.1.6

package/app/utils/state_client.py

@@ -0,0 +1,60 @@
+"""封装 Agent 会话、运行结果与记忆状态的远端操作。"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import httpx
+
+from app.utils.control_auth import ControlAuth
+
+
+class AgentStateError(RuntimeError):
+    pass
+
+
+class AgentStateClient:
+    """访问远端 Agent 状态接口的认证客户端。"""
+
+    def __init__(
+        self,
+        *,
+        http: httpx.AsyncClient,
+        base_url: str,
+        auth: ControlAuth,
+    ) -> None:
+        self._http = http
+        self._base_url = base_url.rstrip("/")
+        self._auth = auth
+
+    @property
+    def headers(self) -> dict[str, str]:
+        return self._auth.headers()
+
+    async def post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        response = await self._http.post(
+            f"{self._base_url}{path}",
+            json=payload,
+            headers=self.headers,
+            timeout=15,
+        )
+        if response.status_code >= 400:
+            raise AgentStateError(f"state request failed: {response.status_code}")
+        body = response.json()
+        if not isinstance(body, dict):
+            raise AgentStateError("state service returned a non-object response")
+        return body
+
+    async def load_messages(self, conversation_id: str) -> list[dict[str, Any]]:
+        body = await self.post(
+            "/agent/getConversationMessages",
+            {"conversationId": conversation_id},
+        )
+        messages = body.get("messages", [])
+        return messages if isinstance(messages, list) else []
+
+    async def complete_run(
+        self,
+        payload: dict[str, Any],
+    ) -> None:
+        await self.post("/agent/completeRun", payload)

package/app/views/__init__.py

@@ -0,0 +1 @@
+"""存放 ASGI 应用的 HTTP 路由与请求边界处理。"""

package/app/views/auth.py

@@ -0,0 +1,189 @@
+"""校验 CAS 与定时任务身份并构造可信运行身份。"""
+
+from __future__ import annotations
+
+import json
+import re
+from base64 import b64decode
+from binascii import Error as Base64Error
+from datetime import datetime, timezone
+from http.cookies import CookieError, SimpleCookie
+from zoneinfo import ZoneInfo
+
+import httpx
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+from starlette.requests import Request
+
+from app.utils.control_auth import ControlAuth
+from app.views.errors import RouteError
+
+_COOKIE_VALUE = re.compile(r"^[A-Za-z0-9._:@+-]{1,256}$")
+_SCHEDULED_TOKEN = re.compile(r"^[A-Za-z0-9_-]{43}$")
+_FORBIDDEN_SCHEDULED_HEADERS = {
+    "x-agent-scheduled-run-id",
+    "x-agent-user-id",
+    "x-agent-org-code",
+    "x-agent-effective-at",
+    "x-agent-effective-timezone",
+}
+
+
+class ScheduledIdentity(BaseModel):
+    model_config = ConfigDict(extra="forbid", populate_by_name=True)
+
+    authentication_type: str = Field(alias="authenticationType")
+    run_id: int = Field(alias="runId", strict=True, gt=0)
+    task_id: int = Field(alias="taskId", strict=True, gt=0)
+    user_id: str = Field(alias="userId", min_length=1, max_length=256)
+    org_code: str = Field(alias="orgCode", min_length=1, max_length=64)
+    role_id: str = Field(alias="roleId", min_length=1, max_length=64)
+    prompt: str = Field(min_length=1, max_length=65536)
+    scheduled_at: datetime = Field(alias="scheduledAt")
+    timezone: str = Field(min_length=1, max_length=128)
+    expires_at: datetime = Field(alias="expiresAt")
+
+
+class _ScheduledEnvelope(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    success: bool
+    data: ScheduledIdentity
+
+
+async def authenticate_user(request: Request, runtime):
+    session_id = _cookies(request.headers.get("cookie", "")).get("CASSESSIONID")
+    if not session_id:
+        raise RouteError(401, "auth error")
+    try:
+        response = await runtime.http.get(
+            f"{runtime.settings.union_base_url}/common/getUserInfo",
+            headers={"Cookie": f"CASSESSIONID={session_id}"},
+            timeout=10,
+        )
+        body = response.json()
+        data = body.get("data", {}) if response.status_code < 400 else {}
+    except (httpx.HTTPError, ValueError):
+        raise RouteError(401, "auth error")
+    if (
+        not isinstance(data, dict)
+        or not isinstance(data.get("userId"), str)
+        or not data["userId"]
+    ):
+        raise RouteError(401, "auth error")
+    permissions = data.get("permissions") or []
+    if isinstance(permissions, str):
+        permissions = [permissions]
+    required = runtime.settings.required_permission
+    if required and required not in permissions:
+        raise RouteError(403, "auth error")
+    return {
+        "authentication_type": "CAS",
+        "user_id": data["userId"],
+        "org_code": data.get("orgCode", ""),
+        "permissions": tuple(permissions),
+        "control_auth": ControlAuth.cas(session_id),
+    }
+
+
+def reject_scheduled_on_cas_endpoint(request: Request) -> None:
+    authorization = request.headers.get("authorization")
+    if not authorization or not authorization.lower().startswith("scheduled "):
+        return
+    if request.headers.get("cookie") is not None:
+        raise RouteError(400, "mixed_authentication")
+    raise RouteError(401, "auth error")
+
+
+def scheduled_token(request: Request) -> str:
+    authorizations = request.headers.getlist("authorization")
+    if len(authorizations) != 1:
+        raise RouteError(401, "scheduled_auth_error")
+    authorization = authorizations[0]
+    cookie_present = request.headers.get("cookie") is not None
+    is_scheduled = bool(
+        authorization and authorization.lower().startswith("scheduled ")
+    )
+    if cookie_present and is_scheduled:
+        raise RouteError(400, "mixed_authentication")
+    if cookie_present or not authorization:
+        raise RouteError(401, "scheduled_auth_error")
+    scheme, separator, token = authorization.partition(" ")
+    if separator != " " or scheme.lower() != "scheduled" or not _valid_token(token):
+        raise RouteError(401, "scheduled_auth_error")
+    return token
+
+
+async def strict_scheduled_body(request: Request) -> None:
+    if _FORBIDDEN_SCHEDULED_HEADERS.intersection(request.headers.keys()):
+        raise RouteError(422, "invalid_scheduled_request")
+    try:
+        payload = json.loads(await request.body())
+    except (json.JSONDecodeError, UnicodeDecodeError):
+        raise RouteError(422, "invalid_scheduled_request")
+    if not isinstance(payload, dict) or payload:
+        raise RouteError(422, "invalid_scheduled_request")
+
+
+async def authenticate_scheduled_run(runtime, token: str) -> ScheduledIdentity:
+    try:
+        response = await runtime.http.post(
+            f"{runtime.settings.union_base_url}/agent/scheduledExecutionIdentity",
+            headers=ControlAuth.scheduled(token).headers(),
+            json={},
+            timeout=10,
+        )
+    except (httpx.TimeoutException, httpx.NetworkError) as exc:
+        raise RouteError(503, "scheduled_identity_unavailable") from exc
+    if response.status_code == 401:
+        raise RouteError(401, "scheduled_auth_error")
+    if response.status_code == 403:
+        raise RouteError(403, "forbidden")
+    if response.status_code >= 500:
+        raise RouteError(503, "scheduled_identity_unavailable")
+    if response.status_code >= 400:
+        raise RouteError(502, "scheduled_identity_unavailable")
+    try:
+        envelope = _ScheduledEnvelope.model_validate(response.json())
+        identity = envelope.data
+        if not envelope.success or identity.authentication_type != "SCHEDULED":
+            raise ValueError
+        if identity.scheduled_at.tzinfo is None or identity.expires_at.tzinfo is None:
+            raise ValueError
+        if identity.expires_at <= datetime.now(timezone.utc):
+            raise ValueError
+        ZoneInfo(identity.timezone)
+        if any(
+            not value.strip()
+            for value in (
+                identity.user_id,
+                identity.org_code,
+                identity.role_id,
+                identity.prompt,
+            )
+        ):
+            raise ValueError
+    except (ValidationError, ValueError, KeyError, json.JSONDecodeError):
+        raise RouteError(502, "scheduled_identity_unavailable")
+    return identity
+
+
+def _valid_token(token: str) -> bool:
+    if not _SCHEDULED_TOKEN.fullmatch(token):
+        return False
+    try:
+        return len(b64decode(token + "=", altchars=b"-_", validate=True)) == 32
+    except (Base64Error, ValueError):
+        return False
+
+
+def _cookies(header: str) -> dict[str, str]:
+    try:
+        cookie = SimpleCookie()
+        cookie.load(header)
+    except CookieError:
+        return {}
+    return {
+        name: morsel.value
+        for name, morsel in cookie.items()
+        if _COOKIE_VALUE.fullmatch(morsel.value)
+    }

package/app/views/errors.py

@@ -0,0 +1,19 @@
+"""定义视图层统一 HTTP 错误及响应格式。"""
+
+from starlette.exceptions import HTTPException
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+
+class RouteError(HTTPException):
+    pass
+
+
+async def http_error(_: Request, exc: RouteError) -> JSONResponse:
+    if isinstance(exc.detail, dict):
+        payload = exc.detail
+    elif isinstance(exc.detail, str):
+        payload = {"error": exc.detail}
+    else:
+        payload = {"detail": exc.detail}
+    return JSONResponse(payload, status_code=exc.status_code, headers=exc.headers)

package/app/views/routes.py

@@ -0,0 +1,25 @@
+"""注册应用对外公开的 HTTP 路由。"""
+
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse
+from starlette.routing import Route
+
+from app.views.errors import RouteError, http_error
+from app.views.streaming_runs import cancel, runs
+from app.views.sync_runs import scheduled_run, sync_run
+
+__all__ = ["RouteError", "build_routes", "http_error"]
+
+
+def build_routes() -> list[Route]:
+    return [
+        Route("/healthcheck.html", healthcheck, methods=["GET"]),
+        Route("/agent/v1/runs", runs, methods=["POST"]),
+        Route("/agent/v1/runs/cancel", cancel, methods=["POST"]),
+        Route("/agent/v1/runs/sync", sync_run, methods=["POST"]),
+        Route("/agent/v1/runs/scheduled", scheduled_run, methods=["POST"]),
+    ]
+
+
+async def healthcheck(_: Request):
+    return PlainTextResponse("success")

package/app/views/run_context.py

@@ -0,0 +1,33 @@
+"""为各类运行构造一致的依赖与状态客户端。"""
+
+import anyio
+
+from app.agent.runtime.deps import RunDeps
+from app.utils.api_client import ApiClient
+from app.utils.state_client import AgentStateClient
+
+
+def build_state_client(runtime, user: dict) -> AgentStateClient:
+    return AgentStateClient(
+        http=runtime.http,
+        base_url=runtime.settings.union_base_url,
+        auth=user["control_auth"],
+    )
+
+
+def build_run_deps(runtime, user: dict, state_client: AgentStateClient) -> RunDeps:
+    return RunDeps(
+        user_id=user["user_id"],
+        permissions=tuple(user.get("permissions") or ()),
+        authentication_type=user.get("authentication_type", "CAS"),
+        org_code=user.get("org_code", ""),
+        scheduled_run_id=user.get("scheduled_run_id"),
+        state_client=state_client,
+        api_client=ApiClient(
+            http=runtime.http,
+            base_url=runtime.settings.union_base_url,
+            auth=user["control_auth"],
+        ),
+        cancelled=anyio.Event(),
+        rag_service=runtime.rag_service,
+    )

package/app/views/streaming_runs.py

@@ -0,0 +1,340 @@
+"""处理 AG-UI 流式运行、取消及消息历史装载。"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import replace
+
+import anyio
+import httpx
+from ag_ui.core import Message
+from pydantic import TypeAdapter, ValidationError
+from pydantic_ai.exceptions import ModelHTTPError
+from pydantic_ai.messages import ModelMessage, ModelRequest, TextContent, UserPromptPart
+from pydantic_ai.ui.ag_ui import AGUIAdapter
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from app.agent.runtime.activity import RunCollector
+from app.agent.runtime.execution import ExecutionPreparationStopped
+from app.agent.runtime.model_errors import model_http_error_details
+from app.utils.state_client import AgentStateError
+from app.views.auth import authenticate_user
+from app.views.errors import RouteError
+from app.views.run_context import build_run_deps, build_state_client
+
+_MESSAGES = TypeAdapter(list[Message])
+_HARNESS_MEMORY_METADATA = "pydantic-ai-harness.memory.v1"
+logger = logging.getLogger(__name__)
+
+
+async def runs(request: Request):
+    runtime = request.app.state.runtime
+    user = await authenticate_user(request, runtime)
+    run_input = _parse_run_input(await request.body())
+    question, user_message = _latest_user_message(run_input.messages)
+    if not question:
+        raise RouteError(422, "a latest user message is required")
+    run_input = run_input.model_copy(update={"messages": [user_message]})
+    state_client = build_state_client(runtime, user)
+    message_history = await _load_message_history(
+        state_client,
+        run_input.thread_id,
+        user_message,
+    )
+    deps = build_run_deps(runtime, user, state_client)
+    collector = RunCollector(
+        root_run_id=run_input.run_id,
+        conversation_id=run_input.thread_id,
+    )
+    collector.root.add_messages(
+        [
+            message.model_dump(by_alias=True, mode="json")
+            for message in run_input.messages
+        ]
+    )
+    record = await _reserve_run(runtime, run_input, deps, collector)
+    root_agent = await _prepare_run(
+        request,
+        runtime,
+        run_input,
+        question,
+        message_history,
+        state_client,
+        deps,
+        collector,
+        record,
+    )
+    record = await _start_run(
+        request,
+        runtime,
+        run_input,
+        root_agent,
+        message_history,
+        deps,
+        collector,
+        record,
+    )
+    return runtime.executions.response(record)
+
+
+async def _reserve_run(runtime, run_input, deps, collector):
+    try:
+        return await runtime.executions.reserve(
+            run_id=run_input.run_id,
+            user_id=deps.user_id,
+            conversation_id=run_input.thread_id,
+            deps=deps,
+            collector=collector,
+        )
+    except ValueError as exc:
+        raise RouteError(409, str(exc)) from exc
+
+
+async def _prepare_run(
+    request,
+    runtime,
+    run_input,
+    question,
+    message_history,
+    state_client,
+    deps,
+    collector,
+    record,
+):
+    preparation_done = anyio.Event()
+
+    async def watch_disconnect() -> None:
+        while not preparation_done.is_set():
+            if await request.is_disconnected():
+                await runtime.executions.cancel(
+                    run_id=run_input.run_id,
+                    user_id=deps.user_id,
+                    conversation_id=run_input.thread_id,
+                )
+                return
+            with anyio.move_on_after(0.1):
+                await preparation_done.wait()
+
+    preparation_phase = "route_resolution"
+    preparation_error: BaseException | None = None
+    try:
+        async with anyio.create_task_group() as task_group:
+            task_group.start_soon(watch_disconnect)
+            try:
+                decision = await runtime.executions.prepare(
+                    record,
+                    lambda: runtime.graph.router.run(
+                        question,
+                        deps=deps,
+                        message_history=message_history,
+                        conversation_id=run_input.thread_id,
+                        usage_limits=runtime.graph.usage_limits,
+                    ),
+                )
+                root_agent = runtime.graph.resolve_root(decision.output)
+                collector.root_agent_name = root_agent.name or "Agent"
+            except BaseException as exc:
+                preparation_error = exc
+            finally:
+                preparation_done.set()
+        if preparation_error is not None:
+            raise preparation_error
+    except ExecutionPreparationStopped as exc:
+        status_code = 504 if exc.error_code == "execution_timeout" else 499
+        raise RouteError(status_code, exc.error_code) from exc
+    except anyio.get_cancelled_exc_class():
+        with anyio.CancelScope(shield=True):
+            await runtime.executions.abandon(
+                record,
+                "cancelled",
+                "request_cancelled",
+            )
+        raise
+    except ModelHTTPError as exc:
+        status, code, message, model_name = model_http_error_details(exc)
+        await runtime.executions.abandon(record, "failed", code)
+        logger.error(
+            "Model request failed phase=%s status=%s model=%s code=%s message=%s",
+            preparation_phase,
+            status,
+            model_name,
+            code,
+            message,
+        )
+        raise RouteError(
+            status,
+            {"error": message, "code": code, "model": model_name},
+        )
+    except Exception:
+        error_code = f"{preparation_phase}_failed"
+        await runtime.executions.abandon(record, "failed", error_code)
+        logger.exception("Run preparation failed phase=%s", preparation_phase)
+        raise RouteError(502, "route resolution failed")
+    return root_agent
+
+
+async def _start_run(
+    request,
+    runtime,
+    run_input,
+    root_agent,
+    message_history,
+    deps,
+    collector,
+    record,
+):
+    try:
+        adapter = AGUIAdapter(
+            agent=root_agent,
+            run_input=run_input,
+            accept=request.headers.get("accept"),
+            allow_uploaded_files=False,
+            preserve_file_data=False,
+        )
+
+        async def capture_root_messages(result) -> None:
+            collector.set_root_messages(
+                [
+                    message.model_dump(by_alias=True, mode="json")
+                    for message in AGUIAdapter.dump_messages(
+                        without_internal_memory(result.new_messages())
+                    )
+                ]
+            )
+
+        stream = adapter.run_stream(
+            output_type=str,
+            deps=deps,
+            message_history=message_history,
+            conversation_id=run_input.thread_id,
+            run_id=run_input.run_id,
+            usage_limits=runtime.graph.usage_limits,
+            on_complete=capture_root_messages,
+        )
+        return await runtime.executions.start(
+            run_id=run_input.run_id,
+            user_id=deps.user_id,
+            conversation_id=run_input.thread_id,
+            adapter=adapter,
+            stream=stream,
+            deps=deps,
+            collector=collector,
+            record=record,
+        )
+    except anyio.get_cancelled_exc_class():
+        with anyio.CancelScope(shield=True):
+            await runtime.executions.abandon(
+                record,
+                "cancelled",
+                "request_cancelled",
+            )
+        raise
+    except ValueError as exc:
+        await runtime.executions.abandon(record, "failed", "run_start_failed")
+        raise RouteError(409, str(exc)) from exc
+    except Exception as exc:
+        await runtime.executions.abandon(record, "failed", "run_start_failed")
+        raise RouteError(500, "run setup failed") from exc
+
+
+def without_internal_memory(messages: list[ModelMessage]) -> list[ModelMessage]:
+    """Remove Harness memory parts while their internal metadata is still intact."""
+    filtered: list[ModelMessage] = []
+    for message in messages:
+        if not isinstance(message, ModelRequest):
+            filtered.append(message)
+            continue
+        parts = []
+        changed = False
+        for part in message.parts:
+            if not isinstance(part, UserPromptPart) or isinstance(part.content, str):
+                parts.append(part)
+                continue
+            content = [
+                item
+                for item in part.content
+                if not (
+                    isinstance(item, TextContent)
+                    and isinstance(item.metadata, str)
+                    and (
+                        item.metadata == _HARNESS_MEMORY_METADATA
+                        or item.metadata.startswith(f"{_HARNESS_MEMORY_METADATA}:")
+                    )
+                )
+            ]
+            if len(content) == len(part.content):
+                parts.append(part)
+            else:
+                changed = True
+                if content:
+                    parts.append(replace(part, content=content))
+        filtered.append(replace(message, parts=parts) if changed else message)
+    return filtered
+
+
+async def cancel(request: Request):
+    runtime = request.app.state.runtime
+    user = await authenticate_user(request, runtime)
+    try:
+        payload = await request.json()
+        run_id = str(payload["runId"])
+        conversation_id = str(payload["conversationId"])
+        cancelled = await runtime.executions.cancel(
+            run_id=run_id,
+            user_id=user["user_id"],
+            conversation_id=conversation_id,
+        )
+    except (KeyError, TypeError, ValueError):
+        return JSONResponse({"error": "invalid cancel request"}, status_code=422)
+    except PermissionError:
+        return JSONResponse({"error": "forbidden"}, status_code=403)
+    return JSONResponse({"success": True, "cancelled": cancelled})
+
+
+def _parse_run_input(body: bytes):
+    try:
+        run_input = AGUIAdapter.build_run_input(body)
+    except ValidationError as exc:
+        raise RouteError(
+            422,
+            exc.errors(include_context=False),
+        ) from exc
+
+    if run_input.tools:
+        raise RouteError(400, "frontend tools are not supported")
+    if run_input.context:
+        raise RouteError(400, "frontend context is not supported")
+    if run_input.state:
+        raise RouteError(400, "frontend state is not supported")
+    if run_input.resume:
+        raise RouteError(400, "frontend resume is not supported")
+    return run_input
+
+
+async def _load_message_history(state_client, thread_id: str, user_message):
+    try:
+        stored = await state_client.load_messages(thread_id)
+        stored_messages = _MESSAGES.validate_python(stored)
+        if stored_messages and stored_messages[-1] == user_message:
+            stored_messages.pop()
+        return AGUIAdapter.load_messages(stored_messages)
+    except (AgentStateError, httpx.HTTPError, ValidationError, TypeError, ValueError) as exc:
+        raise RouteError(503, "conversation history unavailable") from exc
+
+
+def _latest_user_message(messages):
+    for message in reversed(messages):
+        body = message.model_dump(by_alias=True)
+        if body.get("role") == "user":
+            content = body.get("content")
+            if isinstance(content, str):
+                return content.strip(), message
+            if isinstance(content, list):
+                text = "\n".join(
+                    str(item.get("text", ""))
+                    for item in content
+                    if isinstance(item, dict) and item.get("type") == "text"
+                ).strip()
+                return text, message
+    return "", None