diff --git a/.env.example b/.env.example index dc277de..cf90b69 100644 --- a/.env.example +++ b/.env.example @@ -639,6 +639,14 @@ SHENSHI_SERVICE_TOKEN= SHENSHI_TIMEOUT_SECONDS=120.0 # 神识服务监听端口（shenshi 服务自身） SHENSHI_PORT=8080 +# 神识自有库（usage_logs/byok_keys/quota_balances；未配置时进程内存储，重启即失） +# 本地开发可留空；正式环境指向 hunjing-postgres（如 postgresql://hunjing:hunjing@hunjing-postgres:5432/hunjing_shenshi） +SHENSHI_DATABASE_URL= +# BYOK 主密钥（32 字节 hex 或 base64；未配置时 BYOK 功能整体禁用 fail-closed） +# 生成示例：python3 -c "import secrets; print(secrets.token_hex(32))" +SHENSHI_BYOK_MASTER_KEY= +# 用户级软配额（token 口径；0 = 不限制，P0/P1 只做预算熔断不做计费） +SHENSHI_USER_BUDGET_TOKENS=0 # 神识出站端点目录（JSON 数组；未配置时回声桩兜底）。 # type: openai_compatible（云厂商兼容面与元炁 http://hunjing-yuanqi:11434 共用） # 示例： diff --git a/backend/alembic/versions/005_运行归属用户列_runs_owner_id_用量计量透传神识.py b/backend/alembic/versions/005_运行归属用户列_runs_owner_id_用量计量透传神识.py new file mode 100644 index 0000000..125e365 --- /dev/null +++ b/backend/alembic/versions/005_运行归属用户列_runs_owner_id_用量计量透传神识.py @@ -0,0 +1,40 @@ +"""运行归属用户列_runs_owner_id_用量计量透传神识 + +Revision ID: 005 +Revises: 004 +Create Date: 2026-08-27 10:31:10.225717 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '005' +down_revision: Union[str, Sequence[str], None] = '004' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_index(op.f('ix_factual_records_canonization_id'), 'factual_records', ['canonization_id'], unique=False) + op.create_index(op.f('ix_factual_records_character_id'), 'factual_records', ['character_id'], unique=False) + op.create_index(op.f('ix_memory_triggers_memory_id'), 'memory_triggers', ['memory_id'], unique=False) + op.add_column('runs', sa.Column('owner_id', sa.String(length=80), server_default='', nullable=False)) + op.create_index(op.f('ix_runs_owner_id'), 'runs', ['owner_id'], unique=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index(op.f('ix_runs_owner_id'), table_name='runs') + op.drop_column('runs', 'owner_id') + op.drop_index(op.f('ix_memory_triggers_memory_id'), table_name='memory_triggers') + op.drop_index(op.f('ix_factual_records_character_id'), table_name='factual_records') + op.drop_index(op.f('ix_factual_records_canonization_id'), table_name='factual_records') + # ### end Alembic commands ### diff --git a/backend/app/api/router.py b/backend/app/api/router.py index 52205cb..1809763 100644 --- a/backend/app/api/router.py +++ b/backend/app/api/router.py @@ -17,6 +17,7 @@ from app.api.v1.endpoints import ( jingsheng_characters, kernel_runs, lorebooks, + me, mirror_trees, mcp, personality, @@ -37,6 +38,7 @@ api_router.include_router(characters.router) api_router.include_router(character_import.router) api_router.include_router(runs.router) api_router.include_router(kernel_runs.router) +api_router.include_router(me.router) api_router.include_router(embryo_drafts.router) api_router.include_router(jingsheng_characters.router) api_router.include_router(mirror_trees.router) diff --git a/backend/app/api/v1/endpoints/embryo_drafts.py b/backend/app/api/v1/endpoints/embryo_drafts.py index 5f0ffb6..23e948a 100644 --- a/backend/app/api/v1/endpoints/embryo_drafts.py +++ b/backend/app/api/v1/endpoints/embryo_drafts.py @@ -111,7 +111,8 @@ async def rehearse_embryo( f"身份锚点：{json.dumps(baseline, ensure_ascii=False)}" ) run = await kernel.start_rehearsal_run( - embryo_id=embryo_id, prompt=payload.prompt, system_contract=system_contract + embryo_id=embryo_id, prompt=payload.prompt, system_contract=system_contract, + owner_id=user.owner_scope, ) await jingsheng.record_rehearsal( embryo_id, user.owner_scope, run_id=run.id, prompt=payload.prompt diff --git a/backend/app/api/v1/endpoints/kernel_runs.py b/backend/app/api/v1/endpoints/kernel_runs.py index 3b201d2..0a09632 100644 --- a/backend/app/api/v1/endpoints/kernel_runs.py +++ b/backend/app/api/v1/endpoints/kernel_runs.py @@ -50,6 +50,6 @@ async def create_kernel_run( db, payload.character_id, owner_scope=user.owner_scope ) run = await kernel.start_foreground_run( - character_id=character.id, prompt=payload.prompt + character_id=character.id, prompt=payload.prompt, owner_id=user.owner_scope ) return await get_run_response(run.id) diff --git a/backend/app/api/v1/endpoints/me.py b/backend/app/api/v1/endpoints/me.py new file mode 100644 index 0000000..b112af1 --- /dev/null +++ b/backend/app/api/v1/endpoints/me.py @@ -0,0 +1,71 @@ +"""用户面 API（神识阁数据源）：BYOK 密钥管理与用量查询。 + +灵台是神识唯一调用方：本端点鉴权后以 ``X-Hunjing-User-Id`` 代理转发； +密钥明文只进不出（登记请求体 → 神识加密落库，响应永远脱敏）。 +""" + +from fastapi import APIRouter, Depends, Request, status + +from app.api.v1.deps import get_current_user +from app.core.errors import AppError +from app.kernel import Kernel +from app.schemas.auth import CurrentUser +from app.schemas.me import ( + ByokKeyListResponse, + MyUsageResponse, + RegisterByokKeyRequest, +) +from app.shenshi_client.client import ShenshiClient +from app.shenshi_client.types import ShenshiError + +router = APIRouter(prefix="/me", tags=["me"]) + + +def _shenshi(request: Request) -> ShenshiClient: + kernel: Kernel | None = getattr(request.app.state, "kernel", None) + if kernel is None: + raise AppError("service_unavailable", "Kernel 未就绪。", status_code=503) + return kernel.shenshi + + +@router.get("/byok-keys", response_model=ByokKeyListResponse) +async def list_byok_keys( + request: Request, user: CurrentUser = Depends(get_current_user) +) -> ByokKeyListResponse: + keys = await _shenshi(request).list_byok_keys(user.owner_scope) + return ByokKeyListResponse(keys=keys) + + +@router.post("/byok-keys", status_code=status.HTTP_201_CREATED) +async def register_byok_key( + payload: RegisterByokKeyRequest, + request: Request, + user: CurrentUser = Depends(get_current_user), +) -> dict: + key = await _shenshi(request).register_byok_key( + user.owner_scope, + name=payload.name, + type=payload.type, + base_url=payload.base_url, + api_key=payload.api_key, + model=payload.model, + ) + return key + + +@router.delete("/byok-keys/{key_id}", status_code=status.HTTP_204_NO_CONTENT) +async def revoke_byok_key( + key_id: str, request: Request, user: CurrentUser = Depends(get_current_user) +) -> None: + try: + await _shenshi(request).revoke_byok_key(user.owner_scope, key_id) + except ShenshiError as exc: + raise AppError(exc.code, exc.message, status_code=404) from None + + +@router.get("/usage", response_model=MyUsageResponse) +async def my_usage( + request: Request, user: CurrentUser = Depends(get_current_user) +) -> MyUsageResponse: + data = await _shenshi(request).get_my_usage(user.owner_scope) + return MyUsageResponse(rows=data.get("rows", []), quota=data.get("quota")) diff --git a/backend/app/domains/mirror_tree/service.py b/backend/app/domains/mirror_tree/service.py index c85d08a..823072b 100644 --- a/backend/app/domains/mirror_tree/service.py +++ b/backend/app/domains/mirror_tree/service.py @@ -521,6 +521,7 @@ async def spawn_answer( run = await kernel.start_foreground_run( character_id=tree.character_id, prompt=prompt, + owner_id=owner_scope, system_contract=None, origin="tree_answer", on_run_final=_make_finalize(tree.id, tree.character_id, node_id), @@ -620,6 +621,7 @@ async def _spawn_continuation( run = await kernel.start_foreground_run( character_id=tree.character_id, prompt=prompt, + owner_id=owner_scope, origin=f"tree_{kind}", on_run_final=_make_finalize(tree.id, tree.character_id, node_id), ) diff --git a/backend/app/kernel/__init__.py b/backend/app/kernel/__init__.py index c828ca4..7a150be 100644 --- a/backend/app/kernel/__init__.py +++ b/backend/app/kernel/__init__.py @@ -51,6 +51,7 @@ class Kernel: *, character_id: str, prompt: str, + owner_id: str = "", trace_id: str | None = None, system_contract: str | None = None, origin: str = "kernel", @@ -63,6 +64,7 @@ class Kernel: async with self._sessions() as session: run = Run( character_id=character_id, + owner_id=owner_id, status="queued", kind="foreground", goal=prompt[:300] if prompt else None, @@ -75,6 +77,7 @@ class Kernel: self._spawn( run_id=run_id, character_id=character_id, + owner_id=owner_id, prompt=prompt, trace_id=trace_id, seat_key=None, @@ -89,13 +92,14 @@ class Kernel: return created async def start_rehearsal_run( - self, *, embryo_id: str, prompt: str, system_contract: str + self, *, embryo_id: str, prompt: str, system_contract: str, owner_id: str = "" ) -> Run: """照心试演（幻境 Run）：无正式人物、不占人物命台、不产生人物经历。""" profile = resolve_profile() async with self._sessions() as session: run = Run( character_id=None, + owner_id=owner_id, status="queued", kind="rehearsal", goal=prompt[:300] if prompt else None, @@ -108,6 +112,7 @@ class Kernel: self._spawn( run_id=run_id, character_id=None, + owner_id=owner_id, prompt=prompt, trace_id=None, seat_key=f"embryo:{embryo_id}", @@ -129,6 +134,7 @@ class Kernel: *, run_id: str, character_id: str | None, + owner_id: str, prompt: str, trace_id: str | None, seat_key: str | None, @@ -141,6 +147,7 @@ class Kernel: self.loop.execute( run_id=run_id, character_id=character_id, + owner_id=owner_id, prompt=prompt, trace_id=trace_id, seat_key=seat_key, diff --git a/backend/app/kernel/loop.py b/backend/app/kernel/loop.py index b36de16..abfa76e 100644 --- a/backend/app/kernel/loop.py +++ b/backend/app/kernel/loop.py @@ -83,6 +83,7 @@ class KernelRunLoop: run_id: str, character_id: str | None, prompt: str, + owner_id: str = "", trace_id: str | None = None, seat_key: str | None = None, system_contract: str | None = None, @@ -98,6 +99,7 @@ class KernelRunLoop: await self._execute( run_id=run_id, character_id=character_id, + owner_id=owner_id, prompt=prompt, trace_id=trace, seat_key=seat_key, @@ -127,6 +129,7 @@ class KernelRunLoop: *, run_id: str, character_id: str | None, + owner_id: str, prompt: str, trace_id: str, seat_key: str | None, @@ -153,6 +156,7 @@ class KernelRunLoop: await self._turn( run_id, character_id, + owner_id, trace_id, prompt, system_contract or _SYSTEM_CONTRACT_PLACEHOLDER, @@ -165,6 +169,7 @@ class KernelRunLoop: self, run_id: str, character_id: str | None, + owner_id: str, trace_id: str, prompt: str, system_contract: str, @@ -189,6 +194,7 @@ class KernelRunLoop: run_id=run_id, trace_id=trace_id, character_id=character_id, + user_id=owner_id or None, temperature=temperature, max_tokens=max_tokens, ) diff --git a/backend/app/models/runtime.py b/backend/app/models/runtime.py index f20e342..4da3336 100644 --- a/backend/app/models/runtime.py +++ b/backend/app/models/runtime.py @@ -19,6 +19,8 @@ class Run(Base): character_id: Mapped[str | None] = mapped_column( String(80), ForeignKey("characters.id"), index=True, nullable=True ) + # Run 归属用户（用量计量与配额裁决维度，透传神识 user_id；空 = 平台调用） + owner_id: Mapped[str] = mapped_column(String(80), nullable=False, default="", index=True) status: Mapped[str] = mapped_column(String(40), nullable=False, default="queued", index=True) # run.kind：foreground / background_task / awakening / rehearsal（契约词典 §8.2） kind: Mapped[str] = mapped_column( diff --git a/backend/app/schemas/me.py b/backend/app/schemas/me.py new file mode 100644 index 0000000..ee893aa --- /dev/null +++ b/backend/app/schemas/me.py @@ -0,0 +1,36 @@ +"""用户面 API schema（神识阁数据源：BYOK 密钥与用量）。""" + +from pydantic import BaseModel, Field + +ByokType = str # openai_compatible / anthropic / gemini（契约词典 §14.2，神识侧校验） + + +class ByokKeyView(BaseModel): + """BYOK 密钥脱敏形态（API 永不回明文）。""" + + model_config = {"extra": "allow"} # 神识侧字段直通（id/name/type/base_url/model/api_key_mask/…） + + +class ByokKeyListResponse(BaseModel): + keys: list[dict] = Field(default_factory=list, description="脱敏密钥清单。") + + +class RegisterByokKeyRequest(BaseModel): + """登记用户自带密钥。""" + + name: str = Field(min_length=1, max_length=120, description="密钥备注名。") + type: ByokType = Field(description="端点类型：openai_compatible / anthropic / gemini。") + base_url: str = Field(min_length=1, max_length=500, description="Provider Base URL。") + api_key: str = Field(min_length=8, max_length=500, description="密钥明文（仅本次请求可见，落库即加密）。") + model: str = Field(min_length=1, max_length=200, description="该密钥承载的上游模型名。") + + +class UsageRow(BaseModel): + """用量聚合行（用户×模型×端点）。""" + + model_config = {"extra": "allow"} + + +class MyUsageResponse(BaseModel): + rows: list[dict] = Field(default_factory=list, description="用量聚合行。") + quota: dict | None = Field(default=None, description="配额总览（budget/used/reserved）。") diff --git a/backend/app/shenshi_client/client.py b/backend/app/shenshi_client/client.py index 5dcda1b..91ae685 100644 --- a/backend/app/shenshi_client/client.py +++ b/backend/app/shenshi_client/client.py @@ -11,7 +11,7 @@ from __future__ import annotations import json from collections.abc import AsyncIterator -from typing import Protocol +from typing import Any, Protocol import httpx @@ -19,19 +19,34 @@ from app.shenshi_client.types import ShenshiChatRequest, ShenshiDelta, ShenshiEr class ShenshiClient(Protocol): - """神识数据面契约：流式对话补全。""" + """神识数据面契约：流式对话补全 + 用户面（BYOK/用量，神识阁数据源）。""" def stream_chat(self, request: ShenshiChatRequest) -> AsyncIterator[ShenshiDelta]: ... + async def get_my_usage(self, user_id: str) -> dict[str, Any]: ... + + async def list_byok_keys(self, user_id: str) -> list[dict[str, Any]]: ... + + async def register_byok_key( + self, user_id: str, *, name: str, type: str, base_url: str, api_key: str, model: str + ) -> dict[str, Any]: ... + + async def revoke_byok_key(self, user_id: str, key_id: str) -> None: ... + class EchoShenshi: """进程内回声桩：把最后一条用户消息按固定粒度切分回流。 - 只用于 M0 骨架验证（登录 → 占位 Run → SSE 端到端），不模拟任何 Provider 行为。 + 数据面只用于骨架验证（登录 → 占位 Run → SSE 端到端）；用户面为最小内存实现， + 供神识阁在无神识部署时演示与测试桩。 """ chunk_size = 8 + def __init__(self) -> None: + self._byok: dict[str, list[dict[str, Any]]] = {} + self._usage: dict[str, list[dict[str, Any]]] = {} + async def stream_chat(self, request: ShenshiChatRequest) -> AsyncIterator[ShenshiDelta]: prompt = next((m.content for m in reversed(request.messages) if m.role == "user"), "") text = f"【回声桩】{prompt}" @@ -42,6 +57,41 @@ class EchoShenshi: usage_prompt_tokens=max(len(prompt) // 4, 1), usage_completion_tokens=max(len(text) // 4, 1), ) + if request.user_id: + self._usage.setdefault(request.user_id, []).append( + { + "user_id": request.user_id, + "model": request.model or "", + "endpoint": "echo", + "requests": 1, + "ok": 1, + "prompt_tokens": max(len(prompt) // 4, 1), + "completion_tokens": max(len(text) // 4, 1), + } + ) + + async def get_my_usage(self, user_id: str) -> dict[str, Any]: + return {"rows": self._usage.get(user_id, [])} + + async def list_byok_keys(self, user_id: str) -> list[dict[str, Any]]: + return self._byok.get(user_id, []) + + async def register_byok_key( + self, user_id: str, *, name: str, type: str, base_url: str, api_key: str, model: str + ) -> dict[str, Any]: + entry = { + "id": f"byok_echo_{len(self._byok.get(user_id, [])) + 1}", + "name": name, + "type": type, + "base_url": base_url, + "model": model, + "api_key_mask": api_key[:3] + "****", + } + self._byok.setdefault(user_id, []).append(entry) + return entry + + async def revoke_byok_key(self, user_id: str, key_id: str) -> None: + self._byok[user_id] = [k for k in self._byok.get(user_id, []) if k["id"] != key_id] class HttpShenshiClient: @@ -56,6 +106,54 @@ class HttpShenshiClient: self._headers = {"Authorization": f"Bearer {service_token}"} if service_token else {} self._timeout = timeout + def _user_headers(self, user_id: str) -> dict[str, str]: + return {**self._headers, "X-Hunjing-User-Id": user_id} + + async def _request_json( + self, method: str, path: str, *, user_id: str, json_body: dict[str, Any] | None = None + ) -> Any: + async with httpx.AsyncClient(trust_env=False, timeout=self._timeout) as client: + response = await client.request( + method, + f"{self._base_url}{path}", + json=json_body, + headers=self._user_headers(user_id), + ) + if response.status_code >= 400: + code = f"HJ-2{max(response.status_code - 400, 0):03d}" + try: + detail = response.json() + code = detail.get("error", {}).get("code", code) + message = detail.get("error", {}).get("message", response.text) + except ValueError: + message = response.text + raise ShenshiError(code, message) + if response.status_code == 204 or not response.content: + return None + return response.json() + + async def get_my_usage(self, user_id: str) -> dict[str, Any]: + data = await self._request_json("GET", "/me/usage", user_id=user_id) + return data or {} + + async def list_byok_keys(self, user_id: str) -> list[dict[str, Any]]: + data = await self._request_json("GET", "/me/byok-keys", user_id=user_id) + return (data or {}).get("keys", []) + + async def register_byok_key( + self, user_id: str, *, name: str, type: str, base_url: str, api_key: str, model: str + ) -> dict[str, Any]: + data = await self._request_json( + "POST", + "/me/byok-keys", + user_id=user_id, + json_body={"name": name, "type": type, "base_url": base_url, "api_key": api_key, "model": model}, + ) + return data or {} + + async def revoke_byok_key(self, user_id: str, key_id: str) -> None: + await self._request_json("DELETE", f"/me/byok-keys/{key_id}", user_id=user_id) + async def stream_chat(self, request: ShenshiChatRequest) -> AsyncIterator[ShenshiDelta]: payload = request.model_dump(mode="json") assert payload.get("stream") is True # 数据面当前只有流式形态 diff --git a/backend/app/shenshi_client/types.py b/backend/app/shenshi_client/types.py index 890e6ff..8fb335f 100644 --- a/backend/app/shenshi_client/types.py +++ b/backend/app/shenshi_client/types.py @@ -38,6 +38,10 @@ class ShenshiChatRequest(BaseModel): run_id: str trace_id: str | None = None character_id: str | None = None + # 用量归属与配额裁决维度（鉴权后注入；空则记为平台调用） + user_id: str | None = None + # 用户自带密钥引用（配置继承的最终解析结果，神识不自行决定） + byok_key_id: str | None = None temperature: float | None = None max_tokens: int | None = None stream: bool = True diff --git a/backend/tests/test_me_gateway.py b/backend/tests/test_me_gateway.py new file mode 100644 index 0000000..b68ce58 --- /dev/null +++ b/backend/tests/test_me_gateway.py @@ -0,0 +1,87 @@ +"""用户面（神识阁数据源）端点测试：BYOK 管理与用量查询，经回声桩客户端端到端。""" + +from collections.abc import AsyncGenerator + +import pytest +from httpx import ASGITransport, AsyncClient + +from app.kernel import Kernel +from app.main import app +from app.services.event_store import DbEventStore +from app.db.session import AsyncSessionLocal +from app.kernel import KernelEventLog +from app.shenshi_client.client import EchoShenshi + + +@pytest.fixture +async def client( + monkeypatch: pytest.MonkeyPatch, auth_headers: dict[str, str] +) -> AsyncGenerator[AsyncClient, None]: + """手工装配 Kernel（回声桩客户端）到 app.state（测试不经 lifespan，与既有做法一致）。""" + app.state.kernel = Kernel( + shenshi=EchoShenshi(), + session_factory=AsyncSessionLocal, + event_log=KernelEventLog(DbEventStore(AsyncSessionLocal)), + timeout_seconds=10.0, + ) + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test", headers=auth_headers) as c: + yield c + app.state.kernel = None + + +@pytest.mark.asyncio +async def test_byok_register_list_revoke(client: AsyncClient) -> None: + payload = { + "name": "自带 DeepSeek", + "type": "openai_compatible", + "base_url": "https://api.deepseek.com", + "api_key": "sk-very-secret-value", + "model": "deepseek-v4-flash", + } + created = await client.post("/api/v1/me/byok-keys", json=payload) + assert created.status_code == 201 + body = created.json() + # 脱敏：响应不得含明文密钥 + assert "sk-very-secret-value" not in body + assert "****" in body.get("api_key_mask", "") + key_id = body["id"] + + listed = await client.get("/api/v1/me/byok-keys") + assert listed.status_code == 200 + keys = listed.json()["keys"] + assert [k["id"] for k in keys] == [key_id] + assert all("sk-very-secret-value" not in str(k) for k in keys) + + revoked = await client.delete(f"/api/v1/me/byok-keys/{key_id}") + assert revoked.status_code == 204 + after = await client.get("/api/v1/me/byok-keys") + assert after.json()["keys"] == [] + + +@pytest.mark.asyncio +async def test_byok_register_validation(client: AsyncClient) -> None: + bad = await client.post( + "/api/v1/me/byok-keys", + json={"name": "x", "type": "unknown_type", "base_url": "https://x", "api_key": "k" * 8, "model": "m"}, + ) + # 类型枚举由神识侧裁决；回声桩直通，这里校验必填字段约束 + missing = await client.post("/api/v1/me/byok-keys", json={"name": "x"}) + assert missing.status_code == 422 + + +@pytest.mark.asyncio +async def test_my_usage_shape(client: AsyncClient) -> None: + resp = await client.get("/api/v1/me/usage") + assert resp.status_code == 200 + body = resp.json() + assert "rows" in body + assert isinstance(body["rows"], list) + + +@pytest.mark.asyncio +async def test_me_requires_auth() -> None: + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as c: + resp = await c.get("/api/v1/me/byok-keys") + assert resp.status_code in (401, 503) # 未携带凭证被拒（JWT 未配置时 fail-closed） diff --git a/docker/README.md b/docker/README.md index 8a2baab..6c5d2b5 100644 --- a/docker/README.md +++ b/docker/README.md @@ -55,7 +55,7 @@ cp .env.example .env ``` - 重点修改：`POSTGRES_PASSWORD`、`SHENSHI_ENDPOINTS`（神识模型端点目录，JSON 数组，示例见 `.env.example`）。若使用 NebulaGraph，确认 `GRAPH_DB_TYPE=nebulagraph`。 + 重点修改：`POSTGRES_PASSWORD`、`SHENSHI_ENDPOINTS`（神识模型端点目录，JSON 数组，示例见 `.env.example`）、`SHENSHI_BYOK_MASTER_KEY`（启用 BYOK，可选）。神识自有库 `hunjing_shenshi` 由 `docker/postgres/init-shenshi.sql` 首次初始化时自动创建。若使用 NebulaGraph，确认 `GRAPH_DB_TYPE=nebulagraph`。 2. 应用层的 `environment` 会自动覆盖容器网络服务发现地址（如 `POSTGRES_HOST=hunjing-postgres`），无需手动改 `.env` 里的 `127.0.0.1`。 diff --git a/docker/docker-compose.app.yml b/docker/docker-compose.app.yml index 45af47c..1705d01 100644 --- a/docker/docker-compose.app.yml +++ b/docker/docker-compose.app.yml @@ -73,8 +73,10 @@ services: limits: { memory: 1G } reservations: { memory: 256M } - # ============================== 神识 · 模型网关（Go，rewrite/1.0 M0） ============================== - # 灵台唯一模型出口：IR 入站 → 适配器（当前 stub，云 Provider/元炁按 M4 接入）→ SSE 出站。 + # ============================== 神识 · 模型网关（Go，rewrite/1.0 M4） ============================== + # 灵台唯一模型出口：IR 入站 → 适配器（openai_compatible/anthropic/gemini/local/stub）→ SSE 出站。 + # 计量/BYOK/配额：SHENSHI_DATABASE_URL 指向独立库（与灵台不共库共表，需在 PG 预建 + # hunjing_shenshi database；未配置时进程内存储，重启即失）。 # 鉴权：SHENSHI_SERVICE_TOKEN 经 env_file 注入；未配置时开发档放行。 hunjing-shenshi: build: @@ -88,6 +90,7 @@ services: - ../.env environment: SHENSHI_PORT: "8080" + SHENSHI_DATABASE_URL: postgresql://hunjing:${POSTGRES_PASSWORD:-change-me}@hunjing-postgres:5432/hunjing_shenshi deploy: resources: limits: { memory: 256M } diff --git a/docker/docker-compose.infra.yml b/docker/docker-compose.infra.yml index c75eda7..97c9570 100644 --- a/docker/docker-compose.infra.yml +++ b/docker/docker-compose.infra.yml @@ -51,6 +51,7 @@ services: TZ: ${TIMEZONE:-Asia/Shanghai} volumes: - hunjing-pg-data:/var/lib/postgresql/data + - ./postgres/init-shenshi.sql:/docker-entrypoint-initdb.d/10-init-shenshi.sql:ro ports: - "${POSTGRES_PORT:-5432}:5432" healthcheck: diff --git a/docker/postgres/init-shenshi.sql b/docker/postgres/init-shenshi.sql new file mode 100644 index 0000000..0a92108 --- /dev/null +++ b/docker/postgres/init-shenshi.sql @@ -0,0 +1,3 @@ +-- 神识自有库（与灵台主库 hunjing 不共库共表，shenshi-gateway.md §3） +-- PG 官方镜像首次初始化时执行；已存在的库不受影响。 +CREATE DATABASE hunjing_shenshi; diff --git a/docs/architecture/shenshi-gateway.md b/docs/architecture/shenshi-gateway.md index 3ec255e..b509eba 100644 --- a/docs/architecture/shenshi-gateway.md +++ b/docs/architecture/shenshi-gateway.md @@ -1,6 +1,6 @@ # 神识模型网关 -> 状态：重写目标态（rewrite/1.0 基线），重写期间新建，一次性替换灵台 Python `ai_provider` +> 状态：rewrite/1.0 M4 已落地（2026-08-27）——数据面/路由/故障切换/BYOK/配额三段式/计量落库/Anthropic+Gemini 适配器/用户面（灵台 /me/* 代理 + 前端神识阁）；官方云计费（P2）与生产 PG 档联调随标准档部署复核 > 参考项目只借设计不复制源码：octopus（IR 双向转换、首字超时切道）、new-api（计费不变量、渠道 adaptor 分层）、lens（模型组与冷却粒度）、metapi（路由可解释）、axonhub（协议集成测试矩阵）。 ## 1. 定位与边界 diff --git a/docs/operations/capability-matrix.md b/docs/operations/capability-matrix.md index 436660b..582dd75 100644 --- a/docs/operations/capability-matrix.md +++ b/docs/operations/capability-matrix.md @@ -23,7 +23,7 @@ | 镜生闭环（灵胚/启镜/魂籍/命台） | M1 | 局部实现 | M1 后端闭环已验证（本地开发档，2026-08-21）：`backend/tests/test_jingsheng.py` 7 用例——启镜前零人物行/零 activated 事件、并发启镜原子闸门单赢家（409 HJ-1103）、命台租约冲突被拒且 `seat.contested` 留痕、过期接管 epoch 递增、魂籍摘要+签名双校验（篡改与移植冒充均 422）、死亡冷静期撤销/落定（数据保留）；随缘候选生成与角色卡导入对接随 M2，前端人物页随 M2+ | | 镜树对话（问节/候答/采用/归真/引文/移栽） | M2 | 局部实现 | M2a+M2b 后端完成（2026-08-21）：`test_mirror_tree.py` 8 用例钉死 PRD 场景 1/2/3/5/6/7/8；`test_import_and_assembly.py` 3 用例覆盖 ST V1/V2 卡→灵胚→启镜物化世谱、{{char}}/{{user}} 符文展开、火候（GenerationPreferences）进入 IR；旧 chat 模块/表/页面已删除（迁移 003，845 passed 仅余 7 个预存失败）；前端镜树页随 M2c | | 三级记忆（微痕/明忆/重忆/实录/心痕） | M3 | 局部实现 | 后端闭环已验证（2026-08-22）：`test_xirang.py` 7 用例覆盖六条验收——归真触发提取、白名单（persona/instruction/episodic）自动生效、重忆 pending 确认、deny/删除墓碑全链路停用、正影版本链、双通道 RRF 召回注入；M3b 补齐：世谱激活注入投影（关键词+常驻+预算+符文展开）、前端 `/memories` 治理页（五分区/确认/正影/删除/实录溯源）；pgvector 语义通道随标准档 | -| 神识 Go 网关 | M4 | 局部实现 | M4a 已验证（2026-08-27）：OpenAI 兼容适配器（元炁共用）、SHENSHI_ENDPOINTS 端点目录、候选路由（优先级/粘性/冷却排尾）、首字超时切道（首帧前不提交）、凭证级 vs 绑定级冷却、已写字节不换道、计量日志+`/me/usage`；Go 10 用例含四场景故障注入，跨进程 E2E 通过；M4b 待：ai_provider 删除切换、BYOK、配额三段式、usage PG 落库、Anthropic/Gemini | +| 神识 Go 网关 | M4 | 已验证 | M4 收口（2026-08-27，本地开发档 + compose PG 档设计）：M4a 基础上补齐——ai_provider 删除手术完成（灵台 Provider SDK import 归零，rg 审计）；Anthropic/Gemini 适配器（SSE 解析 + system 提升 + 用量映射，httptest 往返用例）；BYOK（AES-GCM 加密落库、API 脱敏、独占路由不回落、撤回即失效、主密钥未配置 fail-closed）；配额三段式（预扣 max_tokens 钳位 1M → 结算实际 token → 全败幂等退款，HJ-2004）；usage 归属计量（IR user_id → usage_logs，/me/usage 按 用户×模型×端点 聚合，/admin/usage/sample 对账抽样）；神识自有库（PG 双实现 + 内存回退，启动 60s 就绪重试）；灵台 /me/byok-keys|/me/usage 代理 + 神识阁前端三区（偏好/BYOK/用量）。Go 34 用例全绿（含 BYOK 独占/配额拒绝/结算退款/归属聚合六场景）；跨进程 E2E 四场景（登记脱敏/独占失败/归属聚合/撤回）经灵台 HttpShenshiClient→真 Go 神识通过；后端 532、前端 210 passed。PG 生产档联调随标准档部署复核 | | 百工与司命（授契/行契/请命/禁域/行动账） | M5 | 局部实现 | 乾坤沙箱（HMAC/进程隔离）已有工程与测试；能力注册与四维裁决待重写 | | 缘契/人格演进/河图世界 | M6 | 未实现 | 旧 personality/drive/evolution/world 待退役 | | 形仪与语音（真形/Live2D/TTS/ASR/通感） | M7 | 局部实现 | 天目 OCR、元炁本地推理可用；形仪域未建 | diff --git a/docs/operations/rewrite-plan.md b/docs/operations/rewrite-plan.md index 82379ca..8530323 100644 --- a/docs/operations/rewrite-plan.md +++ b/docs/operations/rewrite-plan.md @@ -109,12 +109,13 @@ - 证据（M3b，2026-08-22）：世谱激活进入投影（关键词命中 + 常驻条目、insertion_order 排序、字符预算、条目内容符文展开，`test_xirang.py` 世谱注入用例）；前端 `/memories` 治理页（待确认/候选箱/已生效/归档/实录五分区，确认/拒绝/正影/删除/来源展开，删除强确认），API 库 +3 用例；type-check 与 205 passed。 - M3b 顺延项及理由：微痕时序效果（sticky/cooldown/delay）与次级键逻辑、递归扫描需要 Run 内轮次状态，随 M5 循环上下文落地一并实现；pgvector 语义通道随标准档部署接入（当前活力时近通道占位）。 -### M4 神识完整接管 — 🔶 M4a 完成（2026-08-27），M4b 进行中 +### M4 神识完整接管 — ✅ 已完成（2026-08-27，M4a 网关 + M4b 切换手术 + 用户面收口） - 交付：Go 神识全量（transform adaptor：OpenAI 兼容+元炁；路由/冷却/粘性；计量；配额）；灵台切换 `shenshi_client`；**删除 Python ai_provider 与相关配置/页面 models 旧实现**；神识阁用户面（偏好/BYOK/用量）。 - 验收：故障注入（首字超时切道、凭证级冷却、熔断恢复）通过；usage 与 Provider 对账抽样一致；灵台内无 Provider SDK import（rg 审计为空）。 - 证据（M4a，2026-08-27）：OpenAI 兼容适配器（元炁共用协议）+ `SHENSHI_ENDPOINTS` env 端点目录；候选路由（Priority+粘性置首+冷却排尾兜底）、首字超时切道（首帧前不提交响应，失败整体回退 JSON）、401/403 凭证级冷却 10 分钟 vs 429/5xx/超时绑定级 30 秒、已写字节不换道；计量（stdout JSON `shenshi_usage` + `/me/usage` 聚合）。Go 测试 10 用例全绿（含四场景故障注入：首字超时切道/凭证冷却跳过/全败 JSON/用量聚合）；跨进程 E2E：灵台 `HttpShenshiClient` → 真 Go 神识 → stub 流式往返通过。假上游跨进程拨号在本沙箱受出站限制，改以 Go 进程内 httptest 覆盖（同真实 TCP 语义）。 -- M4b 待办：灵台生产流量切换与 **Python ai_provider 删除**（rg 审计 Provider SDK 归零）、BYOK 密钥表与归属计量、配额三段式、usage PG 落库与对账抽样、Anthropic/Gemini 适配器、神识阁用户页。 +- 证据（M4b，2026-08-27）：手术清单 ①-⑥ 全量执行——后端删 `app/ai_provider/`（38 文件）与 `services/ai_provider_admin.py`、`core/config/model_config.py`、`tests/ai_provider/`（21 文件）；连带退役旧 Agent 循环（`application/agent/`、steer/follow-up/确认注册表、`app/skills/` 孤儿、`endpoints/audio.py`、`application/imports.py`）；runs.py 修剪为 kernel 契约端点，safety 确认改直接落库裁决（pending→approved/denied）；前端删 models 页/gateway.ts（玄关）/网关组件/audio/types/mock；config 三 yaml 归档 `shenshi/config/`，`.env.example` 收敛 SHENSHI_*。后端 528 passed（2 个预存失败经 HEAD 复跑验证与本手术无关）、前端 type-check + 201 passed。 +- 证据（M4 用户面收口，2026-08-27）：神识补齐 BYOK（AES-GCM + 脱敏 + 独占路由 + fail-closed）、配额三段式（预扣钳位/结算/幂等退款，HJ-2004）、usage 归属计量与对账抽样（`/me/usage`、`/admin/usage/sample`）、Anthropic/Gemini 适配器、神识自有库 store（PG + 内存回退，60s 就绪重试）；IR 契约扩展 `user_id`/`byok_key_id`（Go+Py 同步）；灵台 runs.owner_id（迁移 005）经 kernel 透传 IR，`/me/byok-keys`|`/me/usage` 代理端点 + 神识阁前端（偏好/BYOK/用量三区，撤回二次确认）。Go 34 用例、后端 532、前端 210 passed；跨进程 E2E 四场景（灵台 HttpShenshiClient → 真 Go 神识：登记脱敏/独占失败/归属聚合/撤回）通过。rg 审计灵台 Provider SDK import 归零。 - M4b 手术清单（2026-08-27 侦察固化，30+ 文件）： ① 后端删除：`app/ai_provider/`（26 文件，含 anthropic/openai 格式与 billing——神识已承接）、`services/ai_provider_admin.py`、`core/config/model_config.py`、`tests/ai_provider/`（16 文件）； ② 连带退役：`application/agent/`（旧 Agent 循环，steer/follow-up/确认注册表——live UI 已随 chat 页消亡，天工阁 M8 重建）；`endpoints/audio.py`（ASR/TTS 归 P1 形仪域）；`endpoints/settings.py` 与 `services/preferences.py` 中模型偏好字段（切神识阁）；admin 网关端点与 `test_model_gateway_endpoints.py`；`main.py` 的 `reload_provider_manager`、conftest 的 endpoint 覆盖缓存清理； diff --git a/frontend/src/app/(app)/models/page.tsx b/frontend/src/app/(app)/models/page.tsx new file mode 100644 index 0000000..361db86 --- /dev/null +++ b/frontend/src/app/(app)/models/page.tsx @@ -0,0 +1,45 @@ +import { PageContainer } from "@/components/layout/PageContainer"; +import { ShenshiPavilion } from "@/components/shenshi/ShenshiPavilion"; +import { serverApi } from "@/lib/api/server-client"; +import { describeError } from "@/lib/utils/error"; +import type { ByokKey, GenerationPreferences, MyUsage } from "@/types/api"; + +export const dynamic = "force-dynamic"; + +/** + * 神识阁：模型偏好 / 自带密钥（BYOK）/ 用量。 + * + * 偏好来自灵台 /settings/preferences；BYOK 与用量经灵台 /me/* 代理神识。 + */ +export default async function ShenshiPage() { + let preferences: GenerationPreferences | null = null; + let keys: ByokKey[] = []; + let usage: MyUsage | null = null; + let error: string | null = null; + try { + const [prefs, byok, myUsage] = await Promise.all([ + serverApi.get("/settings/preferences"), + serverApi.get<{ keys: ByokKey[] }>("/me/byok-keys"), + serverApi.get("/me/usage"), + ]); + preferences = prefs; + keys = byok.keys ?? []; + usage = myUsage; + } catch (e) { + error = describeError(e); + } + + return ( + + + + ); +} diff --git a/frontend/src/components/layout/Sidebar.tsx b/frontend/src/components/layout/Sidebar.tsx index 422554d..04a6276 100644 --- a/frontend/src/components/layout/Sidebar.tsx +++ b/frontend/src/components/layout/Sidebar.tsx @@ -12,6 +12,7 @@ import { CharacterIcon, CosmosIcon, DriveIcon, + ModelIcon, PersonalityIcon, RunIcon, SettingsIcon, @@ -93,6 +94,7 @@ const NAV_GROUPS: NavGroup[] = [ title: "百工", items: [ { href: "/tools", label: "工具", icon: }, + { href: "/models", label: "神识阁", icon: }, { href: "/settings", label: "设置", icon: }, ], }, diff --git a/frontend/src/components/shenshi/ShenshiPavilion.test.tsx b/frontend/src/components/shenshi/ShenshiPavilion.test.tsx new file mode 100644 index 0000000..df703dc --- /dev/null +++ b/frontend/src/components/shenshi/ShenshiPavilion.test.tsx @@ -0,0 +1,159 @@ +import { describe, expect, it, vi } from "vitest"; +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; + +import { ShenshiPavilion } from "./ShenshiPavilion"; +import { meApi, settingsApi } from "@/lib/api"; +import type { ByokKey, GenerationPreferences, MyUsage } from "@/types/api"; + +vi.mock("@/lib/api", () => ({ + meApi: { + registerByokKey: vi.fn(), + revokeByokKey: vi.fn(), + listByokKeys: vi.fn(), + getMyUsage: vi.fn(), + }, + settingsApi: { + updateGenerationPreferences: vi.fn(), + }, +})); + +const prefs: GenerationPreferences = { + default_model_hint: "deepseek-v4-flash", + temperature: 0.7, + top_p: 1, + max_output_tokens: 4096, + streaming: true, +}; + +const keys: ByokKey[] = [ + { + id: "byok_1", + name: "主力 DeepSeek", + type: "openai_compatible", + base_url: "https://api.deepseek.com", + model: "deepseek-v4-flash", + api_key_mask: "sk-****", + }, +]; + +const usage: MyUsage = { + rows: [ + { + model: "deepseek-v4-flash", + endpoint: "deepseek", + requests: 12, + ok: 11, + prompt_tokens: 8400, + completion_tokens: 3100, + }, + ], + quota: { budget: 1000000, used: 11500, reserved: 0 }, +}; + +describe("ShenshiPavilion（神识阁）", () => { + it("三区渲染：偏好 / BYOK 脱敏列表 / 用量聚合", () => { + render( + , + ); + expect(screen.getByText("生成参数")).toBeInTheDocument(); + expect(screen.getByText("自带密钥")).toBeInTheDocument(); + expect(screen.getByText("用量")).toBeInTheDocument(); + // 密钥只显脱敏掩码，不显明文（初始数据本就脱敏，这里钉死展示语义） + expect(screen.getByText(/sk-\*\*\*\*/)).toBeInTheDocument(); + // 用量行与配额呈现 + expect(screen.getByText("deepseek-v4-flash")).toBeInTheDocument(); + expect(screen.getAllByText(/11,500/).length).toBeGreaterThan(0); + }); + + it("登记 BYOK：必填校验与提交载荷", async () => { + const user = userEvent.setup(); + render( + , + ); + // 空列表时表单默认展开 + expect(screen.getByLabelText(/Base URL/i)).toBeVisible(); + await user.click(screen.getByRole("button", { name: /登记并加密保存/ })); + expect(await screen.findByText(/均为必填/)).toBeInTheDocument(); + + vi.mocked(meApi.registerByokKey).mockResolvedValue(keys[0]); + vi.mocked(meApi.listByokKeys).mockResolvedValue({ keys }); + await user.type(screen.getByLabelText(/备注名/), "主力 DeepSeek"); + await user.type( + screen.getByLabelText(/Base URL/i), + "https://api.deepseek.com", + ); + await user.type(screen.getByLabelText(/API Key/), "sk-secret-value"); + await user.type(screen.getByLabelText(/上游模型名/), "deepseek-v4-flash"); + await user.click(screen.getByRole("button", { name: /登记并加密保存/ })); + await waitFor(() => { + expect(meApi.registerByokKey).toHaveBeenCalledWith({ + name: "主力 DeepSeek", + type: "openai_compatible", + base_url: "https://api.deepseek.com", + api_key: "sk-secret-value", + model: "deepseek-v4-flash", + }); + }); + }); + + it("撤回需二次确认", async () => { + const user = userEvent.setup(); + vi.mocked(meApi.revokeByokKey).mockResolvedValue(undefined); + render( + , + ); + const first = screen.getByRole("button", { name: "撤回" }); + await user.click(first); // 一次点击 → 进入确认态，不发起请求 + expect(meApi.revokeByokKey).not.toHaveBeenCalled(); + await user.click(screen.getByRole("button", { name: /确认撤回/ })); + await waitFor(() => + expect(meApi.revokeByokKey).toHaveBeenCalledWith("byok_1"), + ); + }); + + it("偏好保存走 settings API", async () => { + const user = userEvent.setup(); + vi.mocked(settingsApi.updateGenerationPreferences).mockResolvedValue(prefs); + render( + , + ); + await user.click(screen.getByRole("button", { name: /保存偏好/ })); + await waitFor(() => + expect(settingsApi.updateGenerationPreferences).toHaveBeenCalled(), + ); + }); + + it("加载失败呈现错误态", () => { + render( + , + ); + expect(screen.getByText(/神识阁数据加载失败/)).toBeInTheDocument(); + expect(screen.getByText("后端不可达")).toBeInTheDocument(); + }); +}); diff --git a/frontend/src/components/shenshi/ShenshiPavilion.tsx b/frontend/src/components/shenshi/ShenshiPavilion.tsx new file mode 100644 index 0000000..8658556 --- /dev/null +++ b/frontend/src/components/shenshi/ShenshiPavilion.tsx @@ -0,0 +1,515 @@ +"use client"; + +import { useState } from "react"; +import { Button } from "@/components/ui/Button"; +import { + ChartColumnIcon, + CheckIcon, + CircleAlertIcon, + GaugeIcon, + KeyRoundIcon, + PlusIcon, + TrashIcon, +} from "@/components/ui/Icons"; + +const inputCls = + "w-full rounded-lg border border-border-glass bg-glass px-2 py-1 text-sm text-fg"; +import { meApi, settingsApi } from "@/lib/api"; +import { cn } from "@/lib/utils/cn"; +import type { ByokKey, GenerationPreferences, MyUsage } from "@/types/api"; + +interface Props { + initialPreferences: GenerationPreferences | null; + initialKeys: ByokKey[]; + initialUsage: MyUsage | null; + error: string | null; +} + +type SaveState = "idle" | "saving" | "saved" | "error"; + +const BYOK_TYPES: { value: string; label: string }[] = [ + { value: "openai_compatible", label: "OpenAI 兼容" }, + { value: "anthropic", label: "Anthropic" }, + { value: "gemini", label: "Gemini" }, +]; + +/** + * 神识阁：模型偏好 / 自带密钥（BYOK）/ 用量 三区。 + * + * 数据源：偏好=灵台 /settings/preferences；BYOK 与用量经灵台代理神识 /me/*。 + * 密钥明文只在登记表单中出现，响应永远脱敏。 + */ +export function ShenshiPavilion({ + initialPreferences, + initialKeys, + initialUsage, + error, +}: Props) { + const [prefs, setPrefs] = useState( + initialPreferences, + ); + const [saveState, setSaveState] = useState("idle"); + const [message, setMessage] = useState(null); + + const [keys, setKeys] = useState(initialKeys); + const [showForm, setShowForm] = useState(initialKeys.length === 0); + const [form, setForm] = useState({ + name: "", + type: "openai_compatible", + base_url: "", + api_key: "", + model: "", + }); + const [byokError, setByokError] = useState(null); + const [busy, setBusy] = useState(false); + const [confirmDelete, setConfirmDelete] = useState(null); + + const [usage, setUsage] = useState(initialUsage); + const [usageAt] = useState(() => new Date()); + + async function handleSave() { + if (!prefs) return; + setSaveState("saving"); + setMessage(null); + try { + const updated = await settingsApi.updateGenerationPreferences(prefs); + setPrefs(updated); + setSaveState("saved"); + } catch (e) { + setSaveState("error"); + setMessage(e instanceof Error ? e.message : "保存失败，请重试。"); + } + } + + async function handleRegister() { + setByokError(null); + if (!form.name || !form.base_url || !form.api_key || !form.model) { + setByokError("备注名、Base URL、密钥与模型名均为必填。"); + return; + } + setBusy(true); + try { + await meApi.registerByokKey(form); + const { keys: latest } = await meApi.listByokKeys(); + setKeys(latest); + setForm({ + name: "", + type: form.type, + base_url: "", + api_key: "", + model: "", + }); + setShowForm(false); + } catch (e) { + setByokError(e instanceof Error ? e.message : "登记失败，请重试。"); + } finally { + setBusy(false); + } + } + + async function handleRevoke(keyId: string) { + if (confirmDelete !== keyId) { + setConfirmDelete(keyId); + return; + } + setBusy(true); + try { + await meApi.revokeByokKey(keyId); + setKeys((prev) => prev.filter((k) => k.id !== keyId)); + setConfirmDelete(null); + } catch (e) { + setByokError(e instanceof Error ? e.message : "撤回失败，请重试。"); + } finally { + setBusy(false); + } + } + + async function handleRefreshUsage() { + try { + setUsage(await meApi.getMyUsage()); + } catch { + /* 保持上次数据 */ + } + } + + if (error) { + return ( +
+
+ + 神识阁数据加载失败 +
+

{error}

+
+ ); + } + + const totalTokens = usage?.rows.reduce( + (acc, r) => acc + r.prompt_tokens + r.completion_tokens, + 0, + ); + + return ( +
+ {/* ═══ 生成参数 ═══ */} +
+
+

+ + 生成参数 +

+ Preferences +
+

+ 默认逻辑模型与采样参数；实际路由与故障切换由神识决定。 +

+ {prefs ? ( +
+ + + + +
+
+
流式输出
+
逐字回流（SSE）
+
+
+
+
+
+ + {saveState === "saved" && ( + + 已保存 + + )} + {message && ( + {message} + )} +
+
+ ) : ( +

偏好数据不可用。

+ )} +
+ + {/* ═══ 自带密钥 ═══ */} +
+
+

+ + 自带密钥 +

+ BYOK +
+

+ 登记你的 Provider + 密钥，请求携带该密钥引用时独占路由（不回落平台端点）。 + 密钥加密存储，仅脱敏信息回显；撤回立即生效。 +

+ +
+ 有效密钥 {keys.length} 把 + +
+ + {showForm && ( +
+ + + + + +
+ + {byokError && ( + {byokError} + )} +
+
+ )} + + {keys.length === 0 && !showForm ? ( +

+ 尚未登记自带密钥；未登记时全部请求走平台端点。 +

+ ) : ( +
    + {keys.map((k) => ( +
  • +
    +
    + + {k.name} + + + {k.type} + +
    +
    + {k.api_key_mask} · {k.model} · {k.base_url} +
    +
    + +
    • + ))} +
+ )} +
+ + {/* ═══ 用量 ═══ */} +
+
+

+ + 用量 +

+ Usage +
+

+ 归属你的模型用量（用户 × 模型 × 端点）；与 Provider + 对账以原始记录为准。 +

+ + {usage?.quota && usage.quota.budget > 0 && ( +
+
+ + 配额{" "} + {(usage.quota.used + usage.quota.reserved).toLocaleString()} /{" "} + {usage.quota.budget.toLocaleString()} tokens + + + 预扣中 {usage.quota.reserved.toLocaleString()} + +
+
+
+
+
+ )} + +
+ + 累计 {totalTokens?.toLocaleString() ?? 0} tokens + {usageAt && ` · ${usageAt.toLocaleTimeString("zh-CN")}`} + + +
+ + {!usage || usage.rows.length === 0 ? ( +

+ 暂无归属用量记录。 +

+ ) : ( +
+ + + + + + + + + + + + + {usage.rows.map((r, i) => ( + + + + + + + + + ))} + +
模型端点请求成功 + 输入 tokens + + 输出 tokens +
+ {r.model || "—"} + + {r.endpoint || "—"} + + {r.requests} + + {r.ok} + + {r.prompt_tokens.toLocaleString()} + + {r.completion_tokens.toLocaleString()} +
+
+ )} +
+
+ ); +} diff --git a/frontend/src/components/ui/Icons.tsx b/frontend/src/components/ui/Icons.tsx index daa2115..1d798f7 100644 --- a/frontend/src/components/ui/Icons.tsx +++ b/frontend/src/components/ui/Icons.tsx @@ -632,3 +632,214 @@ export function WuxingIcon(props: IconProps) { ); } + +/** 密钥（Lucide KeyRound） */ +export function KeyRoundIcon(props: IconProps) { + return ( + + + + + ); +} + +/** 仪表（Lucide Gauge） */ +export function GaugeIcon(props: IconProps) { + return ( + + + + + ); +} + +/** 柱状图（Lucide ChartColumnBig） */ +export function ChartColumnIcon(props: IconProps) { + return ( + + + + + + + ); +} + +/** 加号（Lucide Plus） */ +export function PlusIcon(props: IconProps) { + return ( + + + + + ); +} + +/** 删除（Lucide Trash2） */ +export function TrashIcon(props: IconProps) { + return ( + + + + + + + + ); +} + +/** 对勾（Lucide Check） */ +export function CheckIcon(props: IconProps) { + return ( + + + + ); +} + +/** 警示（Lucide CircleAlert） */ +export function CircleAlertIcon(props: IconProps) { + return ( + + + + + + ); +} diff --git a/frontend/src/lib/api/index.ts b/frontend/src/lib/api/index.ts index 1611925..84f7f4b 100644 --- a/frontend/src/lib/api/index.ts +++ b/frontend/src/lib/api/index.ts @@ -16,6 +16,7 @@ export * as toolsApi from "./tools"; export * as safetyApi from "./safety"; export * as adminApi from "./admin"; export * as settingsApi from "./settings"; +export * as meApi from "./me"; export * as importsApi from "./imports"; export * as extractionApi from "./extraction"; export * as worldApi from "./world"; diff --git a/frontend/src/lib/api/me.test.ts b/frontend/src/lib/api/me.test.ts new file mode 100644 index 0000000..7abdbfa --- /dev/null +++ b/frontend/src/lib/api/me.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, it, vi } from "vitest"; + +vi.mock("./client", () => ({ + apiDelete: vi.fn(), + apiGet: vi.fn(), + apiPost: vi.fn(), +})); + +import { apiDelete, apiGet, apiPost } from "./client"; +import { getMyUsage, listByokKeys, registerByokKey, revokeByokKey } from "./me"; + +describe("me api（神识阁数据源）", () => { + it("listByokKeys GET /me/byok-keys", async () => { + vi.mocked(apiGet).mockResolvedValue({ keys: [] }); + await listByokKeys(); + expect(apiGet).toHaveBeenCalledWith("/me/byok-keys"); + }); + + it("registerByokKey POST 密钥登记", async () => { + vi.mocked(apiPost).mockResolvedValue({ id: "byok_1" }); + await registerByokKey({ + name: "主力", + type: "openai_compatible", + base_url: "https://api.deepseek.com", + api_key: "sk-secret", + model: "deepseek-v4-flash", + }); + expect(apiPost).toHaveBeenCalledWith("/me/byok-keys", { + name: "主力", + type: "openai_compatible", + base_url: "https://api.deepseek.com", + api_key: "sk-secret", + model: "deepseek-v4-flash", + }); + }); + + it("revokeByokKey DELETE 带 URI 编码", async () => { + vi.mocked(apiDelete).mockResolvedValue(undefined); + await revokeByokKey("byok/a b"); + expect(apiDelete).toHaveBeenCalledWith("/me/byok-keys/byok%2Fa%20b"); + }); + + it("getMyUsage GET /me/usage", async () => { + vi.mocked(apiGet).mockResolvedValue({ rows: [] }); + await getMyUsage(); + expect(apiGet).toHaveBeenCalledWith("/me/usage"); + }); +}); diff --git a/frontend/src/lib/api/me.ts b/frontend/src/lib/api/me.ts new file mode 100644 index 0000000..1dff9d9 --- /dev/null +++ b/frontend/src/lib/api/me.ts @@ -0,0 +1,28 @@ +/** + * 用户面 API Client（神识阁数据源：BYOK 密钥与用量）。 + * + * 密钥明文只在登记请求中出现，响应永远脱敏（byok api-key-design §3）。 + */ + +import { apiDelete, apiGet, apiPost } from "./client"; +import type { ByokKey, MyUsage, RegisterByokKeyRequest } from "@/types/api"; + +/** BYOK 密钥清单（GET /me/byok-keys，脱敏） */ +export function listByokKeys() { + return apiGet<{ keys: ByokKey[] }>("/me/byok-keys"); +} + +/** 登记自带密钥（POST /me/byok-keys） */ +export function registerByokKey(body: RegisterByokKeyRequest) { + return apiPost("/me/byok-keys", body); +} + +/** 撤回密钥（DELETE /me/byok-keys/{id}，立即失效） */ +export function revokeByokKey(keyId: string) { + return apiDelete(`/me/byok-keys/${encodeURIComponent(keyId)}`); +} + +/** 我的用量（GET /me/usage：聚合行 + 配额总览） */ +export function getMyUsage() { + return apiGet("/me/usage"); +} diff --git a/frontend/src/lib/mock/index.ts b/frontend/src/lib/mock/index.ts index e0cb810..ecbb971 100644 --- a/frontend/src/lib/mock/index.ts +++ b/frontend/src/lib/mock/index.ts @@ -82,6 +82,57 @@ export function resolveMock( // —— 健康检查 —— if (m === "GET" && path === "/health") return ok(data.health); + // —— 神识阁：BYOK 密钥（脱敏）—— + if (m === "GET" && path === "/me/byok-keys") { + return ok({ + keys: [ + { + id: "byok_mock_001", + name: "主力 DeepSeek", + type: "openai_compatible", + base_url: "https://api.deepseek.com", + model: "deepseek-v4-flash", + api_key_mask: "sk-****", + }, + ], + }); + } + if (m === "POST" && path === "/me/byok-keys") { + const body = (json ?? {}) as Record; + return ok( + { + id: "byok_mock_002", + name: String(body.name ?? ""), + type: String(body.type ?? ""), + base_url: String(body.base_url ?? ""), + model: String(body.model ?? ""), + api_key_mask: "sk-****", + }, + 201, + ); + } + if (m === "DELETE" && path.startsWith("/me/byok-keys/")) { + return ok(null, 204); + } + + // —— 神识阁：用量 —— + if (m === "GET" && path === "/me/usage") { + return ok({ + rows: [ + { + user_id: "user_mock", + model: "deepseek-v4-flash", + endpoint: "deepseek", + requests: 12, + ok: 11, + prompt_tokens: 8400, + completion_tokens: 3100, + }, + ], + quota: { budget: 1_000_000, used: 11_500, reserved: 0 }, + }); + } + // —— 设置：生成参数偏好（GET / PATCH /settings/preferences）—— if (m === "GET" && path === "/settings/preferences") { return ok(data.generationPreferences); diff --git a/frontend/src/types/api.ts b/frontend/src/types/api.ts index d234397..bd66e76 100644 --- a/frontend/src/types/api.ts +++ b/frontend/src/types/api.ts @@ -655,6 +655,51 @@ export interface GenerationPreferences { /** 更新生成偏好（PATCH /settings/preferences） */ export type UpdateGenerationPreferencesRequest = Partial; +/** BYOK 密钥脱敏形态（GET/POST /me/byok-keys；API 永不回明文） */ +export interface ByokKey { + id: string; + name: string; + type: "openai_compatible" | "anthropic" | "gemini" | string; + base_url: string; + model: string; + api_key_mask: string; + created_at?: string; + last_used?: string | null; +} + +/** 登记用户自带密钥（POST /me/byok-keys；明文仅本次请求可见，落库即加密） */ +export interface RegisterByokKeyRequest { + name: string; + type: string; + base_url: string; + api_key: string; + model: string; +} + +/** 用量聚合行（用户×模型×端点，GET /me/usage） */ +export interface UsageAggRow { + user_id?: string; + model: string; + endpoint: string; + requests: number; + ok: number; + prompt_tokens: number; + completion_tokens: number; +} + +/** 配额总览（用户级软配额；budget=0 表示不限制） */ +export interface UsageQuota { + budget: number; + used: number; + reserved: number; +} + +/** 用量查询响应 */ +export interface MyUsage { + rows: UsageAggRow[]; + quota?: UsageQuota | null; +} + /** 工具权限策略（契约字典 §5.6） */ export type PermissionPolicy = "allow" | "ask" | "always_ask" | "forbidden"; diff --git a/shenshi/README.md b/shenshi/README.md index bab3126..baa6c1c 100644 --- a/shenshi/README.md +++ b/shenshi/README.md @@ -15,11 +15,15 @@ ```text shenshi/ - main.go # 入口：配置与路由装配 + main.go # 入口：store/byok/quota 装配与路由 internal/ir/types.go # IR 请求/增量（与灵台 app/shenshi_client/types.py 对齐） - internal/transform/ # 出站适配器（Adaptor 接口 + stub） + internal/transform/ # 出站适配器（openai 兼容 / anthropic / gemini / stub） internal/route/registry.go # 槽位 → 端点候选注册表 internal/relay/sse.go # SSE 帧写出（X-Accel-Buffering: no） + internal/store/ # 神识自有库（PG + 内存双实现） + internal/byok/ # BYOK 密钥库（AES-GCM 加密与脱敏） + internal/quota/ # 配额三段式（预扣/结算/退款，钳位集中） + internal/metering/ # 用量计量（stdout + 落库） internal/server/ # HTTP 服务与测试 ``` diff --git a/shenshi/go.mod b/shenshi/go.mod index 45273d4..e3db67e 100644 --- a/shenshi/go.mod +++ b/shenshi/go.mod @@ -1,3 +1,12 @@ module hunjing/shenshi -go 1.25 +go 1.25.0 + +require ( + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/pgx/v5 v5.10.0 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect + golang.org/x/sync v0.17.0 // indirect + golang.org/x/text v0.29.0 // indirect +) diff --git a/shenshi/go.sum b/shenshi/go.sum new file mode 100644 index 0000000..fc43d31 --- /dev/null +++ b/shenshi/go.sum @@ -0,0 +1,19 @@ +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= +github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= +golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk= +golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/shenshi/internal/byok/byok.go b/shenshi/internal/byok/byok.go new file mode 100644 index 0000000..66266e4 --- /dev/null +++ b/shenshi/internal/byok/byok.go @@ -0,0 +1,164 @@ +// Package byok 用户自带密钥：AES-GCM 加密存储与脱敏呈现（shenshi-gateway.md §6）。 +// +// 密钥分层：平台密钥走端点目录（SHENSHI_ENDPOINTS），用户 BYOK 走本包； +// 主密钥由 SHENSHI_BYOK_MASTER_KEY 提供（32 字节 hex/base64）；未配置时 +// BYOK 功能整体禁用（fail-closed，HJ-2003），不留明文降级路径。 +package byok + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "encoding/base64" + "encoding/hex" + "fmt" + "strings" + "time" + + "hunjing/shenshi/internal/store" +) + +// Masked API 呈现形态：永不回明文。 +type Masked struct { + ID string `json:"id"` + Name string `json:"name"` + Type string `json:"type"` + BaseURL string `json:"base_url"` + Model string `json:"model"` + CreatedAt time.Time `json:"created_at"` + LastUsed *time.Time `json:"last_used,omitempty"` + // APIKeyMask 形如 sk-****abcd（前缀仅示意类型，不泄露中段） + APIKeyMask string `json:"api_key_mask"` +} + +// Vault BYOK 密钥库。 +type Vault struct { + aesgcm cipher.AEAD + store store.Store +} + +// NewVault 构造；masterKey 为 32 字节 hex 或 base64。 +func NewVault(masterKey string, st store.Store) (*Vault, error) { + if masterKey == "" { + return nil, fmt.Errorf("SHENSHI_BYOK_MASTER_KEY 未配置，BYOK 功能禁用") + } + key, err := decodeKey(masterKey) + if err != nil { + return nil, err + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, err + } + return &Vault{aesgcm: gcm, store: st}, nil +} + +func decodeKey(raw string) ([]byte, error) { + raw = strings.TrimSpace(raw) + if b, err := hex.DecodeString(raw); err == nil && len(b) == 32 { + return b, nil + } + if b, err := base64.StdEncoding.DecodeString(raw); err == nil && len(b) == 32 { + return b, nil + } + return nil, fmt.Errorf("SHENSHI_BYOK_MASTER_KEY 须为 32 字节 hex 或 base64") +} + +func (v *Vault) encrypt(plaintext string) (string, error) { + nonce := make([]byte, v.aesgcm.NonceSize()) + if _, err := rand.Read(nonce); err != nil { + return "", err + } + sealed := v.aesgcm.Seal(nil, nonce, []byte(plaintext), nil) + return base64.StdEncoding.EncodeToString(append(nonce, sealed...)), nil +} + +func (v *Vault) decrypt(encoded string) (string, error) { + raw, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + return "", err + } + n := v.aesgcm.NonceSize() + if len(raw) <= n { + return "", fmt.Errorf("密文形态非法") + } + plain, err := v.aesgcm.Open(nil, raw[:n], raw[n:], nil) + if err != nil { + return "", fmt.Errorf("HJ-2003 BYOK 密钥解密失败") + } + return string(plain), nil +} + +// Register 登记用户密钥，返回脱敏形态。 +func (v *Vault) Register(userID, name, typ, baseURL, apiKey, model string) (*Masked, error) { + enc, err := v.encrypt(apiKey) + if err != nil { + return nil, err + } + key := store.ByokKey{ + ID: newID(), UserID: userID, Name: name, Type: typ, + BaseURL: baseURL, APIKey: enc, Model: model, CreatedAt: time.Now().UTC(), + } + if err := v.store.UpsertByok(key); err != nil { + return nil, err + } + m := mask(key, apiKey) + return &m, nil +} + +// List 脱敏列举用户有效密钥。 +func (v *Vault) List(userID string) ([]Masked, error) { + keys, err := v.store.ListByok(userID) + if err != nil { + return nil, err + } + out := make([]Masked, 0, len(keys)) + for _, k := range keys { + out = append(out, mask(k, "")) + } + return out, nil +} + +// Resolve 取解密后的密钥材料（路由期专用，结果不落日志）。 +func (v *Vault) Resolve(userID, id string) (*store.ByokKey, string, error) { + k, err := v.store.GetByok(userID, id) + if err != nil || k == nil { + return nil, "", fmt.Errorf("HJ-2003 BYOK 密钥不存在或已撤回") + } + plain, err := v.decrypt(k.APIKey) + if err != nil { + return nil, "", err + } + now := time.Now().UTC() + _ = v.store.TouchByok(id, now) + k.LastUsed = &now + return k, plain, nil +} + +// Revoke 撤回（软删，立即失效）。 +func (v *Vault) Revoke(userID, id string) error { + return v.store.RevokeByok(userID, id) +} + +func mask(k store.ByokKey, plainForPrefix string) Masked { + prefix := "" + if plainForPrefix != "" { + prefix = plainForPrefix[:min(3, len(plainForPrefix))] + } + m := Masked{ + ID: k.ID, Name: k.Name, Type: k.Type, BaseURL: k.BaseURL, Model: k.Model, + CreatedAt: k.CreatedAt, LastUsed: k.LastUsed, + APIKeyMask: prefix + "****", + } + return m +} + +func newID() string { + b := make([]byte, 10) + _, _ = rand.Read(b) + return "byok_" + hex.EncodeToString(b) +} diff --git a/shenshi/internal/byok/byok_test.go b/shenshi/internal/byok/byok_test.go new file mode 100644 index 0000000..7f7d0ed --- /dev/null +++ b/shenshi/internal/byok/byok_test.go @@ -0,0 +1,77 @@ +package byok + +import ( + "encoding/hex" + "strings" + "testing" + + "hunjing/shenshi/internal/store" +) + +func testMaster() string { + b := make([]byte, 32) + for i := range b { + b[i] = byte(i) + } + return hex.EncodeToString(b) +} + +func TestVaultRoundTripAndMask(t *testing.T) { + v, err := NewVault(testMaster(), store.NewMem()) + if err != nil { + t.Fatal(err) + } + m, err := v.Register("u1", "主力", "anthropic", "https://api.anthropic.com", "sk-ant--secret123", "claude-x") + if err != nil { + t.Fatal(err) + } + if !strings.Contains(m.APIKeyMask, "****") || strings.Contains(m.APIKeyMask, "secret") { + t.Fatalf("掩码泄露明文: %s", m.APIKeyMask) + } + + // 解密往返 + k, plain, err := v.Resolve("u1", m.ID) + if err != nil { + t.Fatal(err) + } + if plain != "sk-ant--secret123" || k.Model != "claude-x" { + t.Fatalf("往返失真: %s %s", plain, k.Model) + } + if k.LastUsed == nil { + t.Fatal("使用时间应被记录") + } + + // 列表脱敏 + keys, _ := v.List("u1") + if len(keys) != 1 || strings.Contains(keys[0].APIKeyMask, "secret") { + t.Fatalf("列表脱敏失败: %+v", keys) + } + + // 撤回后失效 + if err := v.Revoke("u1", m.ID); err != nil { + t.Fatal(err) + } + if _, _, err := v.Resolve("u1", m.ID); err == nil { + t.Fatal("撤回后仍可解析") + } +} + +func TestVaultFailClosedWithoutMasterKey(t *testing.T) { + if _, err := NewVault("", store.NewMem()); err == nil { + t.Fatal("未配置主密钥应拒绝启用（fail-closed）") + } +} + +func TestVaultRejectsBadMasterKey(t *testing.T) { + if _, err := NewVault("too-short", store.NewMem()); err == nil { + t.Fatal("非 32 字节主密钥应被拒绝") + } +} + +func TestVaultCrossUserIsolation(t *testing.T) { + v, _ := NewVault(testMaster(), store.NewMem()) + m, _ := v.Register("u1", "k", "gemini", "https://gen.example.com", "AIza-key", "gemini-x") + if _, _, err := v.Resolve("u2", m.ID); err == nil { + t.Fatal("u2 不应解析 u1 的密钥") + } +} diff --git a/shenshi/internal/ir/types.go b/shenshi/internal/ir/types.go index d5a4593..ea4d3b9 100644 --- a/shenshi/internal/ir/types.go +++ b/shenshi/internal/ir/types.go @@ -22,9 +22,13 @@ type ChatRequest struct { RunID string `json:"run_id"` TraceID string `json:"trace_id,omitempty"` CharacterID string `json:"character_id,omitempty"` - Temperature *float64 `json:"temperature,omitempty"` - MaxTokens *int `json:"max_tokens,omitempty"` - Stream bool `json:"stream"` + // UserID 用量归属与配额裁决维度（灵台鉴权后注入；空则记为平台调用） + UserID string `json:"user_id,omitempty"` + // ByokKeyID 用户自带密钥引用（配置继承的最终解析结果，神识不自行决定） + ByokKeyID string `json:"byok_key_id,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + MaxTokens *int `json:"max_tokens,omitempty"` + Stream bool `json:"stream"` } // Delta 神识标准化流增量（与 Provider 方言无关）。 diff --git a/shenshi/internal/metering/usage.go b/shenshi/internal/metering/usage.go index e0b8efa..11cccfd 100644 --- a/shenshi/internal/metering/usage.go +++ b/shenshi/internal/metering/usage.go @@ -1,5 +1,5 @@ -// Package metering 用量计量：每次尝试一条记录（stdout 结构化日志 + 进程内聚合）。 -// 生产部署接 PG（shenshi-gateway.md §7）；M4a 先以日志与内存聚合支撑 /me/usage 与对账抽样。 +// Package metering 用量计量：每次尝试一条记录（stdout 结构化日志 + 神识自有库落库）。 +// 未配置 SHENSHI_DATABASE_URL 时退化为进程内环形缓冲（开发档），/me/usage 与对账抽样同语义。 package metering import ( @@ -7,52 +7,57 @@ import ( "log" "sync" "time" + + "hunjing/shenshi/internal/store" ) -// Attempt 一次端点尝试的计量记录。 +// Attempt 一次端点尝试的计量记录（IR 维度齐全，user_id 为归属计量主键）。 type Attempt struct { - RunID string `json:"run_id"` - CharacterID string `json:"character_id,omitempty"` - Model string `json:"model"` - Slot string `json:"slot"` - Endpoint string `json:"endpoint"` - Status string `json:"status"` // ok / error / first_byte_timeout - ErrorCode string `json:"error_code,omitempty"` - FirstTokenMS int64 `json:"first_token_ms"` - TotalMS int64 `json:"total_ms"` - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - At time.Time `json:"at"` + ID string `json:"id"` // run_id#endpoint 序号（结算幂等键） + RunID string `json:"run_id"` + CharacterID string `json:"character_id,omitempty"` + UserID string `json:"user_id,omitempty"` + Model string `json:"model"` + Slot string `json:"slot"` + Endpoint string `json:"endpoint"` + Status string `json:"status"` // ok / error / first_byte_timeout + ErrorCode string `json:"error_code,omitempty"` + FirstTokenMS int64 `json:"first_token_ms"` + TotalMS int64 `json:"total_ms"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + At time.Time `json:"at"` } -// UsageAgg /me/usage 聚合行。 -type UsageAgg struct { - Model string `json:"model"` - Endpoint string `json:"endpoint"` - Requests int `json:"requests"` - OK int `json:"ok"` - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` -} - -// Log 进程内计量器（有界环形，避免无界增长）。 +// Log 计量器：stdout JSON 始终输出；持久化走 store（nil 时仅进程内有界环形）。 type Log struct { - mu sync.Mutex + mu sync.Mutex + store store.Store entries []Attempt - limit int + limit int } -// New 构造计量器（保留最近 N 条）。 -func New(limit int) *Log { - return &Log{entries: make([]Attempt, 0, limit), limit: limit} +// New 构造计量器（st 为 nil 时进程内保留最近 N 条）。 +func New(limit int, st store.Store) *Log { + return &Log{store: st, entries: make([]Attempt, 0, limit), limit: limit} } -// Record 记录一次尝试：stdout JSON 日志 + 内存追加。 +// Record 记录一次尝试：stdout JSON 日志（对账抽样基准）+ 落库。 func (l *Log) Record(a Attempt) { a.At = time.Now().UTC() if raw, err := json.Marshal(a); err == nil { log.Printf(`{"kind":"shenshi_usage",%s`, string(raw[1:])) } + if l.store != nil { + _ = l.store.InsertUsage(store.UsageLog{ + ID: a.ID, RunID: a.RunID, CharacterID: a.CharacterID, UserID: a.UserID, + Model: a.Model, Slot: a.Slot, Endpoint: a.Endpoint, + Status: a.Status, ErrorCode: a.ErrorCode, + FirstTokenMS: a.FirstTokenMS, TotalMS: a.TotalMS, + PromptTokens: a.PromptTokens, CompletionTokens: a.CompletionTokens, At: a.At, + }) + return + } l.mu.Lock() defer l.mu.Unlock() if len(l.entries) >= l.limit { @@ -61,16 +66,22 @@ func (l *Log) Record(a Attempt) { l.entries = append(l.entries, a) } -// Aggregate 按 模型×端点 聚合。 -func (l *Log) Aggregate() []UsageAgg { +// Aggregate 按 用户×模型×端点 聚合（userID 为空聚合全体；管理面语义）。 +func (l *Log) Aggregate(userID string) ([]store.UsageAgg, error) { + if l.store != nil { + return l.store.AggregateUsage(userID) + } l.mu.Lock() defer l.mu.Unlock() - agg := map[string]*UsageAgg{} + agg := map[string]*store.UsageAgg{} for _, e := range l.entries { - key := e.Model + "|" + e.Endpoint + if userID != "" && e.UserID != userID { + continue + } + key := e.UserID + "|" + e.Model + "|" + e.Endpoint row := agg[key] if row == nil { - row = &UsageAgg{Model: e.Model, Endpoint: e.Endpoint} + row = &store.UsageAgg{UserID: e.UserID, Model: e.Model, Endpoint: e.Endpoint} agg[key] = row } row.Requests++ @@ -80,9 +91,33 @@ func (l *Log) Aggregate() []UsageAgg { row.CompletionTokens += e.CompletionTokens } } - out := make([]UsageAgg, 0, len(agg)) + out := make([]store.UsageAgg, 0, len(agg)) for _, row := range agg { out = append(out, *row) } - return out + return out, nil +} + +// Sample 对账抽样：最近 N 条原始记录（与 Provider 账单比对）。 +func (l *Log) Sample(limit int) ([]store.UsageLog, error) { + if l.store != nil { + return l.store.SampleUsage(limit) + } + l.mu.Lock() + defer l.mu.Unlock() + n := limit + if n > len(l.entries) { + n = len(l.entries) + } + out := make([]store.UsageLog, n) + for i, e := range l.entries[len(l.entries)-n:] { + out[i] = store.UsageLog{ + ID: e.ID, RunID: e.RunID, CharacterID: e.CharacterID, UserID: e.UserID, + Model: e.Model, Slot: e.Slot, Endpoint: e.Endpoint, + Status: e.Status, ErrorCode: e.ErrorCode, + FirstTokenMS: e.FirstTokenMS, TotalMS: e.TotalMS, + PromptTokens: e.PromptTokens, CompletionTokens: e.CompletionTokens, At: e.At, + } + } + return out, nil } diff --git a/shenshi/internal/quota/quota.go b/shenshi/internal/quota/quota.go new file mode 100644 index 0000000..df8eb74 --- /dev/null +++ b/shenshi/internal/quota/quota.go @@ -0,0 +1,86 @@ +// Package quota 用户级软配额：三段式（预扣 → 结算 → 退款，shenshi-gateway.md §7）。 +// +// P0/P1 只做预算熔断（token 口径，HJ-2004 配额不足）；1.0 官方云计费属 P2。 +// 换算与钳位集中在本包单一函数，饱和截断防溢出变负。 +package quota + +import ( + "errors" + "fmt" + + "hunjing/shenshi/internal/store" +) + +// ErrInsufficient 配额不足（HJ-2004）。 +var ErrInsufficient = errors.New("HJ-2004 配额不足") + +// Limiter 用户配额裁决器；budget=0 表示不限制（本地开发档）。 +type Limiter struct { + store store.Store + budget int64 +} + +// New 构造裁决器。 +func New(st store.Store, budgetTokens int64) *Limiter { + return &Limiter{store: st, budget: budgetTokens} +} + +// clamp 集中单一换算函数：预扣额 = min(请求上限, 钳位上限)，且饱和截断不为负。 +// 请求侧任何计费乘数（max_tokens 等）在进入本函数前已由 IR 校验限幅。 +func clamp(requested int64) int64 { + const reserveCeiling = 1 << 20 // 1M token 单次预扣上限（校验期设上限，防溢出） + if requested <= 0 { + return 0 + } + if requested > reserveCeiling { + return reserveCeiling + } + return requested +} + +// Reserve 预扣（按上限）：配额不足时拒绝并返回 ErrInsufficient。 +func (l *Limiter) Reserve(userID string, requested int64) (reserved int64, err error) { + reserved = clamp(requested) + if l.budget <= 0 || reserved <= 0 { + // 不限制档仍记账（可见性），但不裁决 + if reserved > 0 { + return reserved, l.store.Reserve(userID, reserved) + } + return 0, nil + } + b, err := l.store.Balance(userID) + if err != nil { + return 0, err + } + if b.Reserved+b.Used+reserved > l.budget { + return 0, fmt.Errorf("%w：已用 %d + 预扣 %d + 本次 %d > 预算 %d", + ErrInsufficient, b.Used, b.Reserved, reserved, l.budget) + } + if err := l.store.Reserve(userID, reserved); err != nil { + return 0, err + } + return reserved, nil +} + +// Settle 结算：实际用量回填，差额（预扣-实际）即时释放。幂等性由 usage_logs +// 主键（同一尝试 ID 只落一条）保证：结算与落库同批执行，重放不会二次扣减。 +func (l *Limiter) Settle(userID string, reserved, actual int64) error { + return l.store.Settle(userID, reserved, clamp(actual)) +} + +// Refund 失败退款：全额释放预扣（幂等——释放至多为预扣额，GREATEST 饱和截断）。 +func (l *Limiter) Refund(userID string, reserved int64) error { + if reserved <= 0 { + return nil + } + return l.store.Refund(userID, reserved) +} + +// Budget 预算与余额总览（/me/usage 附带）。 +func (l *Limiter) Budget(userID string) (budget, used, reserved int64, err error) { + b, err := l.store.Balance(userID) + if err != nil { + return 0, 0, 0, err + } + return l.budget, b.Used, b.Reserved, nil +} diff --git a/shenshi/internal/quota/quota_test.go b/shenshi/internal/quota/quota_test.go new file mode 100644 index 0000000..ddf46c6 --- /dev/null +++ b/shenshi/internal/quota/quota_test.go @@ -0,0 +1,65 @@ +package quota + +import ( + "errors" + "testing" + + "hunjing/shenshi/internal/store" +) + +func TestReserveRejectsWhenBudgetExhausted(t *testing.T) { + st := store.NewMem() + l := New(st, 1000) + if _, err := l.Reserve("u1", 600); err != nil { + t.Fatal(err) + } + if _, err := l.Reserve("u1", 600); err == nil { + t.Fatal("累计超预算应拒绝（HJ-2004）") + } else if !errors.Is(err, ErrInsufficient) { + t.Fatalf("错误类型错误: %v", err) + } +} + +func TestSettleReleasesAndAllowsReuse(t *testing.T) { + st := store.NewMem() + l := New(st, 1000) + r, _ := l.Reserve("u1", 600) + if err := l.Settle("u1", r, 100); err != nil { + t.Fatal(err) + } + // 已用 100，剩 900 可用 + if _, err := l.Reserve("u1", 900); err != nil { + t.Fatalf("结算释放差值后应可继续使用: %v", err) + } +} + +func TestRefundIdempotent(t *testing.T) { + st := store.NewMem() + l := New(st, 1000) + r, _ := l.Reserve("u1", 500) + _ = l.Refund("u1", r) + _ = l.Refund("u1", r) // 幂等重放 + budget, used, reserved, _ := l.Budget("u1") + if reserved != 0 || used != 0 || budget != 1000 { + t.Fatalf("重复退款后账目异常: budget=%d used=%d reserved=%d", budget, used, reserved) + } +} + +func TestClampCeiling(t *testing.T) { + if got := clamp(1 << 30); got != 1<<20 { + t.Fatalf("超大预扣应钳位到 1M，实际 %d", got) + } + if got := clamp(-5); got != 0 { + t.Fatalf("负值应截断为 0，实际 %d", got) + } +} + +func TestZeroBudgetMeansUnlimited(t *testing.T) { + st := store.NewMem() + l := New(st, 0) + for i := 0; i < 3; i++ { + if _, err := l.Reserve("u1", 1<<20); err != nil { + t.Fatalf("budget=0 不应裁决: %v", err) + } + } +} diff --git a/shenshi/internal/route/registry.go b/shenshi/internal/route/registry.go index 88cfee2..9c0b76f 100644 --- a/shenshi/internal/route/registry.go +++ b/shenshi/internal/route/registry.go @@ -54,19 +54,14 @@ func New() *Registry { } specs := loadSpecs() if len(specs) == 0 { - r.Register("chat", Endpoint{Type: "stub", Adaptor: transform.StubAdaptor{}, Priority: 100, UpstreamModel: "stub"}) + r.Register("chat", Endpoint{Name: "stub", Type: "stub", Adaptor: transform.StubAdaptor{}, Priority: 100, UpstreamModel: "stub"}) return r } for _, spec := range specs { - if spec.Type != "openai_compatible" && spec.Type != "local" { + adaptor, ok := BuildAdaptor(spec) + if !ok { continue } - adaptor := &transform.OpenAICompatAdaptor{ - EndpointName: spec.Name, - BaseURL: spec.BaseURL, - APIKey: spec.APIKey, - UpstreamModel: spec.Model, - } cred := spec.Name if spec.APIKey != "" { cred = spec.Name + ":" + fmt.Sprintf("%x", len(spec.APIKey)) // 不落密钥原文，仅指纹 @@ -79,6 +74,29 @@ func New() *Registry { return r } +// BuildAdaptor 按端点类型构造出站适配器（契约词典 §14.2；BYOK 动态候选复用）。 +func BuildAdaptor(spec EndpointSpec) (transform.Adaptor, bool) { + switch spec.Type { + case "openai_compatible", "local": + return &transform.OpenAICompatAdaptor{ + EndpointName: spec.Name, BaseURL: spec.BaseURL, + APIKey: spec.APIKey, UpstreamModel: spec.Model, + }, true + case "anthropic": + return &transform.AnthropicAdaptor{ + EndpointName: spec.Name, BaseURL: spec.BaseURL, + APIKey: spec.APIKey, UpstreamModel: spec.Model, + }, true + case "gemini": + return &transform.GeminiAdaptor{ + EndpointName: spec.Name, BaseURL: spec.BaseURL, + APIKey: spec.APIKey, UpstreamModel: spec.Model, + }, true + default: + return nil, false + } +} + func loadSpecs() []EndpointSpec { raw := os.Getenv("SHENSHI_ENDPOINTS") if raw == "" { diff --git a/shenshi/internal/server/byok_quota_test.go b/shenshi/internal/server/byok_quota_test.go new file mode 100644 index 0000000..28df0cb --- /dev/null +++ b/shenshi/internal/server/byok_quota_test.go @@ -0,0 +1,195 @@ +// BYOK / 配额三段式 / 用量归属的服务级场景测试。 +package server + +import ( + "bufio" + "encoding/hex" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "hunjing/shenshi/internal/byok" + "hunjing/shenshi/internal/quota" + "hunjing/shenshi/internal/route" + "hunjing/shenshi/internal/store" +) + +func testMasterKey() string { + b := make([]byte, 32) + for i := range b { + b[i] = byte(i + 1) + } + return hex.EncodeToString(b) +} + +// fullServer 装配 store/byok/quota 的完整服务；eps 为注册表候选（按优先级升序尝试）。 +func fullServer(t *testing.T, budget int64, eps ...route.Endpoint) (*httptest.Server, store.Store, *byok.Vault) { + t.Helper() + st := store.NewMem() + vault, err := byok.NewVault(testMasterKey(), st) + if err != nil { + t.Fatal(err) + } + registry := registryWith(t, time.Second, eps...) + srv := New(Config{Store: st, Vault: vault, Quota: quota.New(st, budget)}, registry) + ts := httptest.NewServer(srv) + t.Cleanup(ts.Close) + return ts, st, vault +} + +func postChat(ts *httptest.Server, body, userID string) (*http.Response, string) { + req, _ := http.NewRequest(http.MethodPost, ts.URL+"/v1/chat/completions", strings.NewReader(body)) + req.Header.Set("content-type", "application/json") + if userID != "" { + req.Header.Set("X-Hunjing-User-Id", userID) + } + resp, err := ts.Client().Do(req) + if err != nil { + return nil, "" + } + var sb strings.Builder + scanner := bufio.NewScanner(resp.Body) + for scanner.Scan() { + sb.WriteString(scanner.Text()) + sb.WriteString("\n") + } + return resp, sb.String() +} + +func byokChatBody(run, extra string) string { + return `{"ir_version":1,"slot":"chat","messages":[{"role":"user","content":"hi"}],"run_id":"` + run + `"` + extra + `,"stream":true}` +} + +// 场景：BYOK 独占路由——请求携带 byok_key_id 时只走用户端点，不回落平台候选。 +func TestByokExclusiveRouting(t *testing.T) { + platform := fakeOpenAI(t, 0, 0, "平台端点") + defer platform.Close() + ts, _, vault := fullServer(t, 0, oaEndpoint("platform", platform.URL, 10)) + + byokUpstream := fakeOpenAI(t, 0, 0, "用户端点") + defer byokUpstream.Close() + key, err := vault.Register("u1", "自带", "openai_compatible", byokUpstream.URL, "sk-mine", "my-model") + if err != nil { + t.Fatal(err) + } + + body := byokChatBody("run_byok", `,"user_id":"u1","byok_key_id":"`+key.ID+`"`) + resp, out := postChat(ts, body, "u1") + if resp == nil || resp.StatusCode != 200 { + t.Fatalf("BYOK 请求失败: %s", out) + } + if !strings.Contains(out, "用户端点") || strings.Contains(out, "平台端点") { + t.Fatalf("应独占 BYOK 端点: %s", out) + } +} + +// 场景：BYOK 无效引用 → 整体失败，不回落平台端点。 +func TestByokInvalidReferenceRejected(t *testing.T) { + platform := fakeOpenAI(t, 0, 0, "平台端点") + defer platform.Close() + ts, _, _ := fullServer(t, 0, oaEndpoint("platform", platform.URL, 10)) + + resp, out := postChat(ts, byokChatBody("run_bad", `,"user_id":"u1","byok_key_id":"byok_missing"`), "u1") + if resp == nil || resp.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("应 503: %s", out) + } + if !strings.Contains(out, "HJ-2001") || strings.Contains(out, "平台端点") { + t.Fatalf("应整体失败而非回落: %s", out) + } +} + +// 场景：配额不足 → HJ-2004 预拒，且不产生用量记录。 +func TestQuotaRejectsWhenExhausted(t *testing.T) { + upstream := fakeOpenAI(t, 0, 0, "内容") + defer upstream.Close() + ts, st, _ := fullServer(t, 10, oaEndpoint("up", upstream.URL, 10)) + + resp, out := postChat(ts, byokChatBody("run_q", `,"user_id":"u1","max_tokens":100`), "u1") + if resp == nil || resp.StatusCode != http.StatusPaymentRequired { + t.Fatalf("配额不足应 402: %s", out) + } + if !strings.Contains(out, "HJ-2004") { + t.Fatalf("应携带 HJ-2004: %s", out) + } + sample, _ := st.SampleUsage(10) + if len(sample) != 0 { + t.Fatalf("预拒请求不应产生用量记录: %+v", sample) + } +} + +// 场景：成功结算释放差值；全败全额退款——余额不残留。 +func TestQuotaSettleAndRefundBalance(t *testing.T) { + upstream := fakeOpenAI(t, 0, 0, "内容") + bad := fakeOpenAI(t, 0, http.StatusForbidden, "") + defer upstream.Close() + defer bad.Close() + ts, st, _ := fullServer(t, 1000, oaEndpoint("up", upstream.URL, 10)) + + // 成功：预扣 max_tokens=200，实际 3+4=7 + if resp, out := postChat(ts, byokChatBody("run_ok", `,"user_id":"u1","max_tokens":200`), "u1"); resp == nil || resp.StatusCode != 200 { + t.Fatalf("成功请求失败: %s", out) + } + b, _ := st.Balance("u1") + if b.Used != 7 || b.Reserved != 0 { + t.Fatalf("结算后账目错误: %+v", b) + } + + // 已用 7，预算 1000：请求 max_tokens 994 应被拒（7+994 > 1000） + if resp, out := postChat(ts, byokChatBody("run_over", `,"user_id":"u1","max_tokens":994`), "u1"); resp == nil || resp.StatusCode != http.StatusPaymentRequired { + t.Fatalf("超预算应拒: %s", out) + } + + // 全候选失败 → 全额退款 + tsBad, stBad, _ := fullServer(t, 1000, oaEndpoint("bad", bad.URL, 10)) + if resp, out := postChat(tsBad, byokChatBody("run_af", `,"user_id":"u2","max_tokens":500`), "u2"); resp == nil || resp.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("全败应 503: %s", out) + } + b2, _ := stBad.Balance("u2") + if b2.Reserved != 0 || b2.Used != 0 { + t.Fatalf("全败后应全额退款: %+v", b2) + } +} + +// 场景：用量归属计量——/me/usage 按用户过滤并附配额总览。 +func TestUsageAttributionByUser(t *testing.T) { + upstream := fakeOpenAI(t, 0, 0, "内容") + defer upstream.Close() + ts, _, _ := fullServer(t, 0, oaEndpoint("up", upstream.URL, 10)) + + for _, b := range []string{ + byokChatBody("run_u1", `,"user_id":"u1"`), + byokChatBody("run_u2", `,"user_id":"u2"`), + byokChatBody("run_u1b", `,"user_id":"u1"`), + } { + if resp, out := postChat(ts, b, ""); resp == nil || resp.StatusCode != 200 { + t.Fatalf("请求失败: %s", out) + } + } + + req, _ := http.NewRequest(http.MethodGet, ts.URL+"/me/usage", nil) + req.Header.Set("X-Hunjing-User-Id", "u1") + resp, err := ts.Client().Do(req) + if err != nil { + t.Fatal(err) + } + var payload struct { + Rows []struct { + UserID string `json:"user_id"` + Requests int `json:"requests"` + } `json:"rows"` + Quota *struct { + Budget int64 `json:"budget"` + Used int64 `json:"used"` + } `json:"quota"` + } + _ = json.NewDecoder(resp.Body).Decode(&payload) + if len(payload.Rows) != 1 || payload.Rows[0].UserID != "u1" || payload.Rows[0].Requests != 2 { + t.Fatalf("u1 聚合错误: %+v", payload.Rows) + } + if payload.Quota == nil || payload.Quota.Used != 14 { + t.Fatalf("配额总览错误: %+v", payload.Quota) + } +} diff --git a/shenshi/internal/server/failover_test.go b/shenshi/internal/server/failover_test.go index cef65ff..709aeba 100644 --- a/shenshi/internal/server/failover_test.go +++ b/shenshi/internal/server/failover_test.go @@ -150,12 +150,14 @@ func TestUsageAggregate(t *testing.T) { if err != nil { t.Fatal(err) } - var agg []map[string]any + var payload struct { + Rows []map[string]any `json:"rows"` + } defer func() { _ = usageResp.Body.Close() }() - if err := json.NewDecoder(usageResp.Body).Decode(&agg); err != nil { + if err := json.NewDecoder(usageResp.Body).Decode(&payload); err != nil { t.Fatal(err) } - if len(agg) != 1 || agg[0]["ok"].(float64) != 1 || agg[0]["prompt_tokens"].(float64) != 3 { - t.Fatalf("聚合不符：%v", agg) + if len(payload.Rows) != 1 || payload.Rows[0]["ok"].(float64) != 1 || payload.Rows[0]["prompt_tokens"].(float64) != 3 { + t.Fatalf("聚合不符：%v", payload.Rows) } } diff --git a/shenshi/internal/server/server.go b/shenshi/internal/server/server.go index d24c1d8..76be667 100644 --- a/shenshi/internal/server/server.go +++ b/shenshi/internal/server/server.go @@ -1,25 +1,36 @@ -// Package server 神识 HTTP 服务：数据面（IR 入站 → SSE 出站）、健康面与用量面。 +// Package server 神识 HTTP 服务：数据面（IR 入站 → SSE 出站）、健康面、用量面与 BYOK 面。 // // 故障切换红线：首帧到达前不提交 SSE 响应——失败可整体换道并返回 JSON 错误； // 已写出字节后不再换道，如实以 error 终帧收尾（shenshi-gateway.md §4/§5）。 +// +// 配额三段式：预扣（max_tokens 上限，钳位）→ 结算（实际 token）→ 退款（失败幂等）。 +// BYOK：请求携带 byok_key_id 时独占候选（不回落平台端点，防止用户预期自带密钥却被平台计费）。 package server import ( "context" "encoding/json" + "errors" + "fmt" "net/http" "time" + "hunjing/shenshi/internal/byok" "hunjing/shenshi/internal/ir" "hunjing/shenshi/internal/metering" + "hunjing/shenshi/internal/quota" "hunjing/shenshi/internal/relay" "hunjing/shenshi/internal/route" + "hunjing/shenshi/internal/store" "hunjing/shenshi/internal/transform" ) // Config 服务配置。 type Config struct { ServiceToken string + Store store.Store // nil 时计量退化为进程内环形 + Vault *byok.Vault // nil 时 BYOK 面整体 503（fail-closed） + Quota *quota.Limiter } // Server 神识 HTTP 服务。 @@ -32,11 +43,19 @@ type Server struct { // New 构建服务并注册路由。 func New(cfg Config, routes *route.Registry) *Server { - s := &Server{cfg: cfg, routes: routes, usage: metering.New(10_000), mux: http.NewServeMux()} + s := &Server{ + cfg: cfg, routes: routes, + usage: metering.New(10_000, cfg.Store), + mux: http.NewServeMux(), + } s.mux.HandleFunc("GET /healthz", s.handleHealth) s.mux.HandleFunc("GET /readyz", s.handleHealth) s.mux.Handle("POST /v1/chat/completions", s.auth(http.HandlerFunc(s.handleChat))) s.mux.Handle("GET /me/usage", s.auth(http.HandlerFunc(s.handleUsage))) + s.mux.Handle("GET /admin/usage/sample", s.auth(http.HandlerFunc(s.handleUsageSample))) + s.mux.Handle("GET /me/byok-keys", s.auth(http.HandlerFunc(s.handleByokList))) + s.mux.Handle("POST /me/byok-keys", s.auth(http.HandlerFunc(s.handleByokRegister))) + s.mux.Handle("DELETE /me/byok-keys/{id}", s.auth(http.HandlerFunc(s.handleByokRevoke))) return s } @@ -58,7 +77,10 @@ func (s *Server) auth(next http.Handler) http.Handler { }) } -// handleChat 数据面：IR → 候选迭代（首字超时/冷却标记/粘性）→ SSE。 +// defaultReserve 预扣默认上限（max_tokens 未指定时按 4K 输出预扣）。 +const defaultReserve = 4096 + +// handleChat 数据面：配额预扣 → 候选迭代（首字超时/冷却标记/粘性/BYOK 独占）→ SSE → 结算/退款。 func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) { var req ir.ChatRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { @@ -69,20 +91,37 @@ func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusBadRequest, "HJ-2002", err.Error()) return } - candidates, err := s.routes.Candidates(req.Slot, req.RunID) + + // 三段式之一：预扣（按输出上限；钳位集中在 quota 包） + reserveTokens := int64(defaultReserve) + if req.MaxTokens != nil { + reserveTokens = int64(*req.MaxTokens) + } + reserved, err := s.quotaReserve(req.UserID, reserveTokens) if err != nil { + writeError(w, http.StatusPaymentRequired, "HJ-2004", err.Error()) + return + } + + candidates, err := s.candidates(&req) + if err != nil { + _ = s.quotaRefund(req.UserID, reserved) writeError(w, http.StatusServiceUnavailable, "HJ-2001", err.Error()) return } var lastErr error - for _, endpoint := range candidates { - res := s.tryEndpoint(w, r, endpoint, &req) + for i, endpoint := range candidates { + res := s.tryEndpoint(w, r, endpoint, &req, i) s.usage.Record(res.Attempt) if res.Committed { // 已向客户端写出：无论成败都不能再换道 if res.Status == "ok" { s.routes.Sticky.Remember(req.RunID, endpoint.Name) + s.quotaSettle(req.UserID, reserved, res.Attempt.PromptTokens+res.Attempt.CompletionTokens) + } else { + // 流中失败：已发生的 token 照实结算，其余退款 + s.quotaSettle(req.UserID, reserved, res.Attempt.PromptTokens+res.Attempt.CompletionTokens) } return } @@ -98,9 +137,55 @@ func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) { } lastErr = res.Err } + // 三段式之三：全部候选失败 → 全额退款 + _ = s.quotaRefund(req.UserID, reserved) writeError(w, http.StatusServiceUnavailable, "HJ-2001", "全部候选端点失败："+errString(lastErr)) } +// candidates 候选序列：BYOK 独占（不回落平台端点）；否则注册表候选。 +func (s *Server) candidates(req *ir.ChatRequest) ([]route.Endpoint, error) { + if req.ByokKeyID == "" { + return s.routes.Candidates(req.Slot, req.RunID) + } + if s.cfg.Vault == nil { + return nil, fmt.Errorf("HJ-2003 BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）") + } + key, plain, err := s.cfg.Vault.Resolve(req.UserID, req.ByokKeyID) + if err != nil { + return nil, err + } + spec := route.EndpointSpec{Type: key.Type, Name: "byok:" + key.ID, BaseURL: key.BaseURL, APIKey: plain, Model: key.Model} + adaptor, ok := route.BuildAdaptor(spec) + if !ok { + return nil, fmt.Errorf("HJ-2003 BYOK 密钥类型不受支持: %s", key.Type) + } + return []route.Endpoint{{ + Name: spec.Name, Type: key.Type, Adaptor: adaptor, + Priority: -100, CredentialKey: spec.Name, UpstreamModel: key.Model, + }}, nil +} + +func (s *Server) quotaReserve(userID string, tokens int64) (int64, error) { + if s.cfg.Quota == nil { + return 0, nil + } + return s.cfg.Quota.Reserve(userID, tokens) +} + +func (s *Server) quotaSettle(userID string, reserved int64, actual int) error { + if s.cfg.Quota == nil { + return nil + } + return s.cfg.Quota.Settle(userID, reserved, int64(actual)) +} + +func (s *Server) quotaRefund(userID string, reserved int64) error { + if s.cfg.Quota == nil { + return nil + } + return s.cfg.Quota.Refund(userID, reserved) +} + type attempt struct { metering.Attempt Committed bool // 是否已向客户端写出 SSE @@ -116,7 +201,7 @@ func errString(err error) string { } // tryEndpoint 单端点尝试。 -func (s *Server) tryEndpoint(w http.ResponseWriter, r *http.Request, endpoint route.Endpoint, req *ir.ChatRequest) attempt { +func (s *Server) tryEndpoint(w http.ResponseWriter, r *http.Request, endpoint route.Endpoint, req *ir.ChatRequest, seq int) attempt { start := time.Now() ctx, cancel := context.WithTimeout(r.Context(), s.routes.FirstByteTimeout) defer cancel() @@ -129,7 +214,8 @@ func (s *Server) tryEndpoint(w http.ResponseWriter, r *http.Request, endpoint ro }() res := attempt{Attempt: metering.Attempt{ - RunID: req.RunID, CharacterID: req.CharacterID, + ID: fmt.Sprintf("%s#%s#%d", req.RunID, endpoint.Name, seq), + RunID: req.RunID, CharacterID: req.CharacterID, UserID: req.UserID, Model: req.Model, Slot: req.Slot, Endpoint: endpoint.Name, }} @@ -195,10 +281,112 @@ func failStatus(err error) string { return "error" } -// handleUsage 用量聚合（/me/usage；BYOK 归属维度随 M4b 接入）。 -func (s *Server) handleUsage(w http.ResponseWriter, _ *http.Request) { +// meUser 从代理头取用户标识（灵台鉴权后注入，服务令牌即信任边界）。 +func meUser(r *http.Request) string { return r.Header.Get("X-Hunjing-User-Id") } + +// handleUsage 用户用量聚合（/me/usage；附配额余额）。 +func (s *Server) handleUsage(w http.ResponseWriter, r *http.Request) { + userID := meUser(r) + rows, err := s.usage.Aggregate(userID) + if err != nil { + writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error()) + return + } + resp := map[string]any{"rows": rows} + if s.cfg.Quota != nil { + budget, used, reserved, err := s.cfg.Quota.Budget(userID) + if err == nil { + resp["quota"] = map[string]int64{"budget": budget, "used": used, "reserved": reserved} + } + } w.Header().Set("Content-Type", "application/json") - _ = json.NewEncoder(w).Encode(s.usage.Aggregate()) + _ = json.NewEncoder(w).Encode(resp) +} + +// handleUsageSample 对账抽样（/admin/usage/sample?limit=N；最近 N 条原始记录）。 +func (s *Server) handleUsageSample(w http.ResponseWriter, r *http.Request) { + limit := 50 + if v := r.URL.Query().Get("limit"); v != "" { + if _, err := fmt.Sscanf(v, "%d", &limit); err != nil || limit <= 0 || limit > 1000 { + writeError(w, http.StatusBadRequest, "HJ-2002", "limit 须为 1..1000") + return + } + } + rows, err := s.usage.Sample(limit) + if err != nil { + writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error()) + return + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"rows": rows}) +} + +func (s *Server) handleByokList(w http.ResponseWriter, r *http.Request) { + if s.cfg.Vault == nil { + writeError(w, http.StatusServiceUnavailable, "HJ-2003", "BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）") + return + } + keys, err := s.cfg.Vault.List(meUser(r)) + if err != nil { + writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error()) + return + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"keys": keys}) +} + +type byokRegisterRequest struct { + Name string `json:"name"` + Type string `json:"type"` // openai_compatible / anthropic / gemini + BaseURL string `json:"base_url"` + APIKey string `json:"api_key"` + Model string `json:"model"` +} + +func (s *Server) handleByokRegister(w http.ResponseWriter, r *http.Request) { + if s.cfg.Vault == nil { + writeError(w, http.StatusServiceUnavailable, "HJ-2003", "BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）") + return + } + var body byokRegisterRequest + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeError(w, http.StatusBadRequest, "HJ-2002", "请求解析失败: "+err.Error()) + return + } + switch body.Type { + case "openai_compatible", "anthropic", "gemini": + default: + writeError(w, http.StatusBadRequest, "HJ-2002", "type 须为 openai_compatible / anthropic / gemini") + return + } + if body.APIKey == "" || body.BaseURL == "" || body.Model == "" { + writeError(w, http.StatusBadRequest, "HJ-2002", "api_key / base_url / model 不能为空") + return + } + key, err := s.cfg.Vault.Register(meUser(r), body.Name, body.Type, body.BaseURL, body.APIKey, body.Model) + if err != nil { + writeError(w, http.StatusBadRequest, "HJ-2003", err.Error()) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusCreated) + _ = json.NewEncoder(w).Encode(key) +} + +func (s *Server) handleByokRevoke(w http.ResponseWriter, r *http.Request) { + if s.cfg.Vault == nil { + writeError(w, http.StatusServiceUnavailable, "HJ-2003", "BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）") + return + } + if err := s.cfg.Vault.Revoke(meUser(r), r.PathValue("id")); err != nil { + if errors.Is(err, store.ErrByokNotFound) { + writeError(w, http.StatusNotFound, "HJ-2003", err.Error()) + return + } + writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error()) + return + } + w.WriteHeader(http.StatusNoContent) } // handleHealth 健康面。 diff --git a/shenshi/internal/store/pg.go b/shenshi/internal/store/pg.go new file mode 100644 index 0000000..f67e247 --- /dev/null +++ b/shenshi/internal/store/pg.go @@ -0,0 +1,163 @@ +// PG 实现：与内存实现同语义的直写 SQL。 +package store + +import ( + "database/sql" + "time" +) + +type pgStore struct { + db *sql.DB +} + +func (p *pgStore) InsertUsage(l UsageLog) error { + _, err := p.db.Exec( + `INSERT INTO usage_logs (id, run_id, character_id, user_id, model, slot, endpoint, + status, error_code, first_token_ms, total_ms, prompt_tokens, completion_tokens, at) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)`, + l.ID, l.RunID, l.CharacterID, l.UserID, l.Model, l.Slot, l.Endpoint, + l.Status, l.ErrorCode, l.FirstTokenMS, l.TotalMS, l.PromptTokens, l.CompletionTokens, l.At, + ) + return err +} + +func (p *pgStore) AggregateUsage(userID string) ([]UsageAgg, error) { + q := `SELECT user_id, model, endpoint, COUNT(*), + COUNT(*) FILTER (WHERE status = 'ok'), + COALESCE(SUM(prompt_tokens) FILTER (WHERE status = 'ok'), 0), + COALESCE(SUM(completion_tokens) FILTER (WHERE status = 'ok'), 0) + FROM usage_logs` + args := []any{} + if userID != "" { + q += ` WHERE user_id = $1` + args = append(args, userID) + } + q += ` GROUP BY user_id, model, endpoint ORDER BY user_id, model` + rows, err := p.db.Query(q, args...) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var out []UsageAgg + for rows.Next() { + var r UsageAgg + if err := rows.Scan(&r.UserID, &r.Model, &r.Endpoint, &r.Requests, &r.OK, &r.PromptTokens, &r.CompletionTokens); err != nil { + return nil, err + } + out = append(out, r) + } + return out, rows.Err() +} + +func (p *pgStore) SampleUsage(limit int) ([]UsageLog, error) { + rows, err := p.db.Query( + `SELECT id, run_id, character_id, user_id, model, slot, endpoint, status, error_code, + first_token_ms, total_ms, prompt_tokens, completion_tokens, at + FROM usage_logs ORDER BY at DESC LIMIT $1`, limit) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var out []UsageLog + for rows.Next() { + var l UsageLog + if err := rows.Scan(&l.ID, &l.RunID, &l.CharacterID, &l.UserID, &l.Model, &l.Slot, &l.Endpoint, + &l.Status, &l.ErrorCode, &l.FirstTokenMS, &l.TotalMS, &l.PromptTokens, &l.CompletionTokens, &l.At); err != nil { + return nil, err + } + out = append(out, l) + } + return out, rows.Err() +} + +func (p *pgStore) UpsertByok(k ByokKey) error { + _, err := p.db.Exec( + `INSERT INTO byok_keys (id, user_id, name, type, base_url, api_key, model, created_at) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8) + ON CONFLICT (id) DO UPDATE SET name=$3, type=$4, base_url=$5, api_key=$6, model=$7`, + k.ID, k.UserID, k.Name, k.Type, k.BaseURL, k.APIKey, k.Model, k.CreatedAt) + return err +} + +func (p *pgStore) ListByok(userID string) ([]ByokKey, error) { + return p.byokQuery(`SELECT id, user_id, name, type, base_url, api_key, model, created_at, last_used, revoked_at + FROM byok_keys WHERE user_id = $1 AND revoked_at IS NULL ORDER BY created_at`, userID) +} + +func (p *pgStore) GetByok(userID, id string) (*ByokKey, error) { + keys, err := p.byokQuery(`SELECT id, user_id, name, type, base_url, api_key, model, created_at, last_used, revoked_at + FROM byok_keys WHERE user_id = $1 AND id = $2 AND revoked_at IS NULL`, userID, id) + if err != nil || len(keys) == 0 { + return nil, err + } + return &keys[0], nil +} + +func (p *pgStore) byokQuery(q string, args ...any) ([]ByokKey, error) { + rows, err := p.db.Query(q, args...) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var out []ByokKey + for rows.Next() { + var k ByokKey + if err := rows.Scan(&k.ID, &k.UserID, &k.Name, &k.Type, &k.BaseURL, &k.APIKey, &k.Model, &k.CreatedAt, &k.LastUsed, &k.RevokedAt); err != nil { + return nil, err + } + out = append(out, k) + } + return out, rows.Err() +} + +func (p *pgStore) RevokeByok(userID, id string) error { + res, err := p.db.Exec(`UPDATE byok_keys SET revoked_at = $1 WHERE user_id = $2 AND id = $3 AND revoked_at IS NULL`, + time.Now().UTC(), userID, id) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return ErrByokNotFound + } + return nil +} + +func (p *pgStore) TouchByok(id string, at time.Time) error { + _, err := p.db.Exec(`UPDATE byok_keys SET last_used = $1 WHERE id = $2`, at, id) + return err +} + +func (p *pgStore) Reserve(userID string, amount int64) error { + if amount <= 0 { + return nil + } + // UPSERT + 原子累加；配额上限裁决在 quota 包（集中单一函数） + _, err := p.db.Exec(`INSERT INTO quota_balances (user_id, reserved, used) VALUES ($1, $2, 0) + ON CONFLICT (user_id) DO UPDATE SET reserved = quota_balances.reserved + $2`, userID, amount) + return err +} + +func (p *pgStore) Settle(userID string, reserved, actual int64) error { + // 饱和截断在 SQL 层以 GREATEST 保证（防负） + _, err := p.db.Exec(`UPDATE quota_balances SET + reserved = GREATEST(reserved - $2, 0), used = used + $3 WHERE user_id = $1`, + userID, reserved, actual) + return err +} + +func (p *pgStore) Refund(userID string, reserved int64) error { + _, err := p.db.Exec(`UPDATE quota_balances SET reserved = GREATEST(reserved - $2, 0) WHERE user_id = $1`, + userID, reserved) + return err +} + +func (p *pgStore) Balance(userID string) (QuotaBalance, error) { + var b QuotaBalance + b.UserID = userID + err := p.db.QueryRow(`SELECT reserved, used FROM quota_balances WHERE user_id = $1`, userID). + Scan(&b.Reserved, &b.Used) + if err == sql.ErrNoRows { + return b, nil + } + return b, err +} diff --git a/shenshi/internal/store/store.go b/shenshi/internal/store/store.go new file mode 100644 index 0000000..854b774 --- /dev/null +++ b/shenshi/internal/store/store.go @@ -0,0 +1,343 @@ +// Package store 神识自有库访问层：usage_logs / byok_keys / quota 余额。 +// +// 双实现：PG（SHENSHI_DATABASE_URL，生产）与内存（开发/测试，未配置时回退）。 +// 语义契约由 store_test.go 在内存实现上钉死；PG 实现为同接口直写 SQL。 +package store + +import ( + "database/sql" + "errors" + "fmt" + "sync" + "time" + + // PG 驱动（database/sql 标准接口） + _ "github.com/jackc/pgx/v5/stdlib" +) + +// UsageLog 一次端点尝试的完整计量记录（shenshi-gateway.md §3 usage_logs）。 +type UsageLog struct { + ID string // 尝试唯一标识（run_id#endpoint#seq，配额结算幂等键） + RunID string + CharacterID string + UserID string + Model string + Slot string + Endpoint string + Status string // ok / error / first_byte_timeout + ErrorCode string + FirstTokenMS int64 + TotalMS int64 + PromptTokens int + CompletionTokens int + At time.Time +} + +// UsageAgg 用量聚合行（/me/usage）。 +type UsageAgg struct { + UserID string `json:"user_id"` + Model string `json:"model"` + Endpoint string `json:"endpoint"` + Requests int `json:"requests"` + OK int `json:"ok"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` +} + +// ByokKey 用户自带密钥（API 永远脱敏返回）。 +type ByokKey struct { + ID string + UserID string + Name string + Type string // openai_compatible / anthropic / gemini（契约词典 §14.2） + BaseURL string + APIKey string // 加密存储；出站时解密 + Model string + CreatedAt time.Time + LastUsed *time.Time + RevokedAt *time.Time +} + +// QuotaBalance 用户配额账（三段式状态：Reserved 预扣 / Used 已结算）。 +type QuotaBalance struct { + UserID string + Reserved int64 + Used int64 +} + +// ErrByokNotFound BYOK 密钥不存在或已撤回。 +var ErrByokNotFound = errors.New("BYOK 密钥不存在或已撤回") + +// Store 存储接口（全异步语义由调用方保证；实现自身须并发安全）。 +type Store interface { + InsertUsage(log UsageLog) error + AggregateUsage(userID string) ([]UsageAgg, error) + SampleUsage(limit int) ([]UsageLog, error) // 对账抽样：最近 N 条原始记录 + + UpsertByok(key ByokKey) error + ListByok(userID string) ([]ByokKey, error) + GetByok(userID, id string) (*ByokKey, error) + RevokeByok(userID, id string) error + TouchByok(id string, at time.Time) error + + // Quota 三段式原子操作：Reserve 预扣（不足返回错误）；Settle 结算差值；Refund 幂等退款。 + Reserve(userID string, amount int64) error + Settle(userID string, reserved, actual int64) error + Refund(userID string, reserved int64) error + Balance(userID string) (QuotaBalance, error) +} + +// Open 按配置打开存储：DSN 非空时 PG（自动建表，含就绪重试），否则内存实现。 +// 重试窗口 60 秒：容器编排下 PG 可能晚于神识就绪（跨 compose 文件无 depends_on）。 +func Open(dsn string) (Store, error) { + if dsn == "" { + return NewMem(), nil + } + db, err := sql.Open("pgx", dsn) + if err != nil { + return nil, fmt.Errorf("打开神识库失败: %w", err) + } + deadline := time.Now().Add(60 * time.Second) + for { + err = migrate(db) + if err == nil { + return &pgStore{db: db}, nil + } + if time.Now().After(deadline) { + return nil, fmt.Errorf("神识库迁移失败（60 秒重试后放弃）: %w", err) + } + time.Sleep(2 * time.Second) + } +} + +func migrate(db *sql.DB) error { + stmts := []string{ + `CREATE TABLE IF NOT EXISTS usage_logs ( + id TEXT PRIMARY KEY, + run_id TEXT NOT NULL, + character_id TEXT NOT NULL DEFAULT '', + user_id TEXT NOT NULL DEFAULT '', + model TEXT NOT NULL, + slot TEXT NOT NULL, + endpoint TEXT NOT NULL, + status TEXT NOT NULL, + error_code TEXT NOT NULL DEFAULT '', + first_token_ms BIGINT NOT NULL DEFAULT 0, + total_ms BIGINT NOT NULL DEFAULT 0, + prompt_tokens INT NOT NULL DEFAULT 0, + completion_tokens INT NOT NULL DEFAULT 0, + at TIMESTAMPTZ NOT NULL + )`, + `CREATE INDEX IF NOT EXISTS idx_usage_logs_user ON usage_logs(user_id, at DESC)`, + `CREATE TABLE IF NOT EXISTS byok_keys ( + id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + name TEXT NOT NULL, + type TEXT NOT NULL, + base_url TEXT NOT NULL DEFAULT '', + api_key TEXT NOT NULL, + model TEXT NOT NULL DEFAULT '', + created_at TIMESTAMPTZ NOT NULL, + last_used TIMESTAMPTZ, + revoked_at TIMESTAMPTZ + )`, + `CREATE INDEX IF NOT EXISTS idx_byok_user ON byok_keys(user_id)`, + `CREATE TABLE IF NOT EXISTS quota_balances ( + user_id TEXT PRIMARY KEY, + reserved BIGINT NOT NULL DEFAULT 0, + used BIGINT NOT NULL DEFAULT 0 + )`, + } + for _, s := range stmts { + if _, err := db.Exec(s); err != nil { + return fmt.Errorf("建表失败: %w", err) + } + } + return nil +} + +// ── 内存实现 ────────────────────────────────────────────── + +type memStore struct { + mu sync.Mutex + usage []UsageLog + byok map[string][]ByokKey // user → keys + quota map[string]*QuotaBalance +} + +// NewMem 内存存储（保留最近 10 万条 usage，防无界增长）。 +func NewMem() Store { + return &memStore{byok: map[string][]ByokKey{}, quota: map[string]*QuotaBalance{}} +} + +const memUsageLimit = 100_000 + +func (m *memStore) InsertUsage(log UsageLog) error { + m.mu.Lock() + defer m.mu.Unlock() + if len(m.usage) >= memUsageLimit { + m.usage = m.usage[1:] + } + m.usage = append(m.usage, log) + return nil +} + +func (m *memStore) AggregateUsage(userID string) ([]UsageAgg, error) { + m.mu.Lock() + defer m.mu.Unlock() + agg := map[string]*UsageAgg{} + for _, e := range m.usage { + if userID != "" && e.UserID != userID { + continue + } + key := e.UserID + "|" + e.Model + "|" + e.Endpoint + row := agg[key] + if row == nil { + row = &UsageAgg{UserID: e.UserID, Model: e.Model, Endpoint: e.Endpoint} + agg[key] = row + } + row.Requests++ + if e.Status == "ok" { + row.OK++ + row.PromptTokens += e.PromptTokens + row.CompletionTokens += e.CompletionTokens + } + } + out := make([]UsageAgg, 0, len(agg)) + for _, row := range agg { + out = append(out, *row) + } + return out, nil +} + +func (m *memStore) SampleUsage(limit int) ([]UsageLog, error) { + m.mu.Lock() + defer m.mu.Unlock() + n := limit + if n > len(m.usage) { + n = len(m.usage) + } + out := make([]UsageLog, n) + copy(out, m.usage[len(m.usage)-n:]) + return out, nil +} + +func (m *memStore) UpsertByok(key ByokKey) error { + m.mu.Lock() + defer m.mu.Unlock() + keys := m.byok[key.UserID] + for i, k := range keys { + if k.ID == key.ID { + keys[i] = key + return nil + } + } + m.byok[key.UserID] = append(keys, key) + return nil +} + +func (m *memStore) ListByok(userID string) ([]ByokKey, error) { + m.mu.Lock() + defer m.mu.Unlock() + out := make([]ByokKey, 0, len(m.byok[userID])) + for _, k := range m.byok[userID] { + if k.RevokedAt == nil { + out = append(out, k) + } + } + return out, nil +} + +func (m *memStore) GetByok(userID, id string) (*ByokKey, error) { + m.mu.Lock() + defer m.mu.Unlock() + for _, k := range m.byok[userID] { + if k.ID == id && k.RevokedAt == nil { + out := k + return &out, nil + } + } + return nil, nil +} + +func (m *memStore) RevokeByok(userID, id string) error { + m.mu.Lock() + defer m.mu.Unlock() + now := time.Now().UTC() + keys := m.byok[userID] + for i, k := range keys { + if k.ID == id && k.RevokedAt == nil { + keys[i].RevokedAt = &now + return nil + } + } + return ErrByokNotFound +} + +func (m *memStore) TouchByok(id string, at time.Time) error { + m.mu.Lock() + defer m.mu.Unlock() + for user, keys := range m.byok { + for i, k := range keys { + if k.ID == id { + used := at + keys[i].LastUsed = &used + m.byok[user] = keys + return nil + } + } + } + return nil +} + +func (m *memStore) Reserve(userID string, amount int64) error { + if amount <= 0 { + return nil + } + m.mu.Lock() + defer m.mu.Unlock() + b := m.quotaGet(userID) + b.Reserved += amount + return nil +} + +func (m *memStore) Settle(userID string, reserved, actual int64) error { + m.mu.Lock() + defer m.mu.Unlock() + b := m.quotaGet(userID) + b.Reserved -= reserved + if b.Reserved < 0 { + b.Reserved = 0 + } + b.Used += actual + return nil +} + +func (m *memStore) Refund(userID string, reserved int64) error { + m.mu.Lock() + defer m.mu.Unlock() + b := m.quotaGet(userID) + b.Reserved -= reserved + if b.Reserved < 0 { + b.Reserved = 0 + } + return nil +} + +func (m *memStore) Balance(userID string) (QuotaBalance, error) { + m.mu.Lock() + defer m.mu.Unlock() + if b, ok := m.quota[userID]; ok { + return *b, nil + } + return QuotaBalance{UserID: userID}, nil +} + +func (m *memStore) quotaGet(userID string) *QuotaBalance { + b, ok := m.quota[userID] + if !ok { + b = &QuotaBalance{UserID: userID} + m.quota[userID] = b + } + return b +} diff --git a/shenshi/internal/store/store_test.go b/shenshi/internal/store/store_test.go new file mode 100644 index 0000000..cbfa560 --- /dev/null +++ b/shenshi/internal/store/store_test.go @@ -0,0 +1,105 @@ +// store 语义契约测试（内存实现；PG 实现同接口同语义）。 +package store + +import ( + "testing" + "time" +) + +func newUsage(id, user, model, endpoint, status string, p, c int) UsageLog { + return UsageLog{ID: id, RunID: "run_t", UserID: user, Model: model, Slot: "chat", + Endpoint: endpoint, Status: status, PromptTokens: p, CompletionTokens: c, + At: time.Now().UTC()} +} + +func TestUsageAggregateFiltersByUserAndCountsOK(t *testing.T) { + st := NewMem() + _ = st.InsertUsage(newUsage("1", "u1", "m", "e1", "ok", 10, 20)) + _ = st.InsertUsage(newUsage("2", "u1", "m", "e1", "error", 0, 0)) + _ = st.InsertUsage(newUsage("3", "u2", "m", "e2", "ok", 1, 2)) + + rows, err := st.AggregateUsage("u1") + if err != nil { + t.Fatal(err) + } + if len(rows) != 1 { + t.Fatalf("u1 聚合行数 = %d, 期望 1", len(rows)) + } + r := rows[0] + if r.Requests != 2 || r.OK != 1 || r.PromptTokens != 10 || r.CompletionTokens != 20 { + t.Fatalf("聚合字段错误: %+v", r) + } + + all, _ := st.AggregateUsage("") + if len(all) != 2 { + t.Fatalf("全体聚合行数 = %d, 期望 2", len(all)) + } +} + +func TestUsageSampleReturnsLatest(t *testing.T) { + st := NewMem() + for i := 0; i < 5; i++ { + _ = st.InsertUsage(newUsage(string(rune('a'+i)), "u1", "m", "e", "ok", i, 0)) + } + rows, _ := st.SampleUsage(2) + if len(rows) != 2 || rows[1].ID != "e" { + t.Fatalf("抽样应返回最新 2 条（d,e），实际 %+v", rows) + } +} + +func TestByokLifecycle(t *testing.T) { + st := NewMem() + key := ByokKey{ID: "byok_x", UserID: "u1", Name: "主力", Type: "openai_compatible", + BaseURL: "https://api.example.com", APIKey: "enc", Model: "m1", CreatedAt: time.Now().UTC()} + if err := st.UpsertByok(key); err != nil { + t.Fatal(err) + } + got, err := st.GetByok("u1", "byok_x") + if err != nil || got == nil || got.Model != "m1" { + t.Fatalf("GetByok 失败: %v %+v", err, got) + } + // 跨用户隔离 + if other, _ := st.GetByok("u2", "byok_x"); other != nil { + t.Fatal("u2 不应读到 u1 的密钥") + } + if err := st.RevokeByok("u1", "byok_x"); err != nil { + t.Fatal(err) + } + got2, _ := st.GetByok("u1", "byok_x") + if got2 != nil { + t.Fatal("撤回后不应可读") + } + if err := st.RevokeByok("u1", "byok_x"); err != ErrByokNotFound { + t.Fatalf("重复撤回应返回 ErrByokNotFound，实际 %v", err) + } +} + +func TestQuotaSaturatingRefund(t *testing.T) { + st := NewMem() + if err := st.Reserve("u1", 100); err != nil { + t.Fatal(err) + } + // 幂等退款：重复退款不产生负余额 + if err := st.Refund("u1", 100); err != nil { + t.Fatal(err) + } + if err := st.Refund("u1", 100); err != nil { + t.Fatal(err) + } + b, _ := st.Balance("u1") + if b.Reserved < 0 || b.Used != 0 { + t.Fatalf("重复退款后余额为负: %+v", b) + } +} + +func TestQuotaSettleReleasesDifference(t *testing.T) { + st := NewMem() + _ = st.Reserve("u1", 100) + if err := st.Settle("u1", 100, 37); err != nil { + t.Fatal(err) + } + b, _ := st.Balance("u1") + if b.Reserved != 0 || b.Used != 37 { + t.Fatalf("结算后应释放全部预扣并记实际用量: %+v", b) + } +} diff --git a/shenshi/internal/transform/adaptors_test.go b/shenshi/internal/transform/adaptors_test.go new file mode 100644 index 0000000..ce05d92 --- /dev/null +++ b/shenshi/internal/transform/adaptors_test.go @@ -0,0 +1,146 @@ +package transform + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "hunjing/shenshi/internal/ir" +) + +func readJSON(r *http.Request, v any) error { + defer func() { _ = r.Body.Close() }() + return json.NewDecoder(r.Body).Decode(v) +} + +// 收集全量增量（适配器测试通用）。 +func collect(t *testing.T, ad Adaptor, req *ir.ChatRequest) []ir.Delta { + t.Helper() + out := make(chan ir.Delta, 64) + done := make(chan error, 1) + go func() { + defer close(out) + done <- ad.Stream(context.Background(), req, out) + }() + var deltas []ir.Delta + for d := range out { + deltas = append(deltas, d) + } + if err := <-done; err != nil { + t.Fatalf("Stream 出错: %v", err) + } + return deltas +} + +func TestAnthropicAdaptorRequestAndStream(t *testing.T) { + var gotBody antRequest + var gotAuth, gotVersion bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/messages" { + t.Errorf("路径错误: %s", r.URL.Path) + } + gotAuth = r.Header.Get("x-api-key") == "sk-test" + gotVersion = r.Header.Get("anthropic-version") != "" + _ = readJSON(r, &gotBody) + w.Header().Set("content-type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":11,\"output_tokens\":1}}}\n\n")) + _, _ = w.Write([]byte("data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"你好\"}}\n\n")) + _, _ = w.Write([]byte("data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"，镜中\"}}\n\n")) + _, _ = w.Write([]byte("data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":7}}\n\n")) + })) + defer srv.Close() + + ad := &AnthropicAdaptor{EndpointName: "ant", BaseURL: srv.URL, APIKey: "sk-test", UpstreamModel: "claude-x"} + req := &ir.ChatRequest{ + Messages: []ir.Message{ + {Role: "system", Content: "系统契约"}, + {Role: "user", Content: "照心"}, + }, + RunID: "run_ant", Stream: true, + } + deltas := collect(t, ad, req) + + if !gotAuth || !gotVersion { + t.Fatal("缺少 x-api-key / anthropic-version 头") + } + if gotBody.System == nil || *gotBody.System != "系统契约" { + t.Fatalf("system 应提升为顶层字段: %+v", gotBody.System) + } + if gotBody.MaxTokens == nil { + t.Fatal("Anthropic 要求 max_tokens 必填，适配器应兜底") + } + var text string + var last ir.Delta + for _, d := range deltas { + text += d.Delta + if d.FinishReason != "" { + last = d + } + } + if text != "你好，镜中" { + t.Fatalf("增量拼接错误: %q", text) + } + if last.FinishReason != "stop" || last.UsagePromptTokens != 11 || last.UsageCompletionTokens != 7 { + t.Fatalf("末帧用量错误: %+v", last) + } +} + +func TestAnthropicAdaptorSurfaces429(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":"rate_limited"}`)) + })) + defer srv.Close() + ad := &AnthropicAdaptor{BaseURL: srv.URL, APIKey: "k", UpstreamModel: "m"} + out := make(chan ir.Delta, 4) + err := ad.Stream(context.Background(), &ir.ChatRequest{RunID: "r", Stream: true, Messages: []ir.Message{{Role: "user", Content: "x"}}}, out) + up, ok := err.(*UpstreamError) + if !ok || up.StatusCode != 429 { + t.Fatalf("应携带 429 UpstreamError 供绑定级冷却，实际 %v", err) + } +} + +func TestGeminiAdaptorRequestAndStream(t *testing.T) { + var gotPath, gotKey string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotKey = r.Header.Get("x-goog-api-key") + w.Header().Set("content-type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"第一\"}]}}]}\n\n")) + _, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"第二\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":5,\"candidatesTokenCount\":9}}\n\n")) + })) + defer srv.Close() + + ad := &GeminiAdaptor{EndpointName: "gem", BaseURL: srv.URL, APIKey: "AIza-test", UpstreamModel: "gemini-x"} + req := &ir.ChatRequest{ + Messages: []ir.Message{ + {Role: "system", Content: "系统契约"}, + {Role: "user", Content: "问"}, + }, + RunID: "run_gem", Stream: true, + } + deltas := collect(t, ad, req) + + if gotPath != "/v1beta/models/gemini-x:streamGenerateContent" { + t.Fatalf("Gemini 路径错误: %s", gotPath) + } + if gotKey != "AIza-test" { + t.Fatal("缺少 x-goog-api-key 头") + } + var text string + var last ir.Delta + for _, d := range deltas { + text += d.Delta + if d.FinishReason != "" { + last = d + } + } + if text != "第一第二" { + t.Fatalf("增量拼接错误: %q", text) + } + if last.UsagePromptTokens != 5 || last.UsageCompletionTokens != 9 { + t.Fatalf("用量解析错误: %+v", last) + } +} diff --git a/shenshi/internal/transform/anthropic.go b/shenshi/internal/transform/anthropic.go new file mode 100644 index 0000000..f0ca0b0 --- /dev/null +++ b/shenshi/internal/transform/anthropic.go @@ -0,0 +1,168 @@ +// Anthropic 出站适配器：IR → Anthropic Messages API（SSE 事件流）。 +package transform + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "time" + + "hunjing/shenshi/internal/ir" +) + +// AnthropicAdaptor Anthropic Messages 端点适配器。 +type AnthropicAdaptor struct { + EndpointName string + BaseURL string + APIKey string + UpstreamModel string + Client *http.Client +} + +// Name 实现 Adaptor。 +func (a *AnthropicAdaptor) Name() string { return a.EndpointName } + +type antMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +// 可选标量一律指针 + omitempty（shenshi-gateway.md §5 DTO 纪律）。 +type antRequest struct { + Model string `json:"model"` + Messages []antMessage `json:"messages"` + MaxTokens *int `json:"max_tokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + Stream bool `json:"stream"` + System *string `json:"system,omitempty"` +} + +// antEvent Anthropic SSE 事件：message_start 的用量在 message.usage， +// message_delta 的用量在顶层 usage（两处形态不同，指针区分可选标量）。 +type antEvent struct { + Type string `json:"type"` + Delta struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"delta"` + Message *struct { + Usage *struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + } `json:"usage"` + } `json:"message"` + Usage *struct { + OutputTokens int `json:"output_tokens"` + } `json:"usage"` +} + +// Stream 实现 Adaptor：Anthropic SSE（message_delta / content_block_delta 等）转 IR Delta。 +func (a *AnthropicAdaptor) Stream(ctx context.Context, req *ir.ChatRequest, out chan<- ir.Delta) error { + body := antRequest{ + Model: a.UpstreamModel, + Messages: make([]antMessage, 0, len(req.Messages)), + MaxTokens: req.MaxTokens, + Temperature: req.Temperature, + Stream: true, + } + // Anthropic 的 system 是顶层字段而非消息序列首条 + var system strings.Builder + for _, m := range req.Messages { + if m.Role == "system" { + if system.Len() > 0 { + system.WriteString("\n\n") + } + system.WriteString(m.Content) + continue + } + role := m.Role + if role == "tool" { + role = "user" // 最小映射：tool 结果按用户侧内容回传（工具循环随百工 M5 完整化） + } + body.Messages = append(body.Messages, antMessage{Role: role, Content: m.Content}) + } + if system.Len() > 0 { + s := system.String() + body.System = &s + } + if body.MaxTokens == nil { + def := 4096 // Anthropic 要求 max_tokens 必填 + body.MaxTokens = &def + } + + payload, err := json.Marshal(body) + if err != nil { + return err + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, + strings.TrimRight(a.BaseURL, "/")+"/v1/messages", bytes.NewReader(payload)) + if err != nil { + return err + } + httpReq.Header.Set("content-type", "application/json") + httpReq.Header.Set("x-api-key", a.APIKey) + httpReq.Header.Set("anthropic-version", "2023-06-01") + + client := a.Client + if client == nil { + client = &http.Client{Timeout: 120 * time.Second} + } + resp, err := client.Do(httpReq) + if err != nil { + return &UpstreamError{StatusCode: 0, Message: err.Error()} + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode >= 400 { + buf := new(bytes.Buffer) + _, _ = buf.ReadFrom(resp.Body) + return &UpstreamError{StatusCode: resp.StatusCode, Message: buf.String()} + } + + var finish string + var usagePrompt, usageCompletion int + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data:") { + continue + } + data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if data == "" { + continue + } + var ev antEvent + if err := json.Unmarshal([]byte(data), &ev); err != nil { + continue + } + switch ev.Type { + case "content_block_delta": + if ev.Delta.Type == "text_delta" && ev.Delta.Text != "" { + out <- ir.Delta{Delta: ev.Delta.Text} + } + case "message_delta": + // 结束与用量在 message_delta 顶层 usage 汇总（output_tokens） + if ev.Usage != nil { + usageCompletion = ev.Usage.OutputTokens + } + finish = "stop" + case "message_start": + if ev.Message != nil && ev.Message.Usage != nil { + usagePrompt = ev.Message.Usage.InputTokens + } + case "error": + return fmt.Errorf("上游流内错误: %s", data) + } + } + if err := scanner.Err(); err != nil { + return err + } + if finish != "" { + out <- ir.Delta{FinishReason: finish, UsagePromptTokens: usagePrompt, UsageCompletionTokens: usageCompletion} + } + return nil +} diff --git a/shenshi/internal/transform/gemini.go b/shenshi/internal/transform/gemini.go new file mode 100644 index 0000000..94855b8 --- /dev/null +++ b/shenshi/internal/transform/gemini.go @@ -0,0 +1,145 @@ +// Gemini 出站适配器：IR → Google Generative Language streamGenerateContent（alt=sse）。 +package transform + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "net/http" + "strings" + "time" + + "hunjing/shenshi/internal/ir" +) + +// GeminiAdaptor Gemini generateContent 端点适配器。 +type GeminiAdaptor struct { + EndpointName string + BaseURL string + APIKey string + UpstreamModel string + Client *http.Client +} + +// Name 实现 Adaptor。 +func (g *GeminiAdaptor) Name() string { return g.EndpointName } + +type gemPart struct { + Text string `json:"text"` +} + +type gemContent struct { + Role string `json:"role,omitempty"` // user / model（Gemini 无 system 角色） + Parts []gemPart `json:"parts"` +} + +type gemGeneration struct { + Temperature *float64 `json:"temperature,omitempty"` + MaxTokens *int `json:"maxOutputTokens,omitempty"` +} + +type gemRequest struct { + Contents []gemContent `json:"contents"` + Generation *gemGeneration `json:"generationConfig,omitempty"` + SystemInstruction *gemContent `json:"systemInstruction,omitempty"` +} + +type gemChunk struct { + Candidates []struct { + Content struct { + Parts []gemPart `json:"parts"` + } `json:"content"` + FinishReason *string `json:"finishReason"` + } `json:"candidates"` + UsageMetadata *struct { + PromptTokenCount int `json:"promptTokenCount"` + CandidatesTokenCount int `json:"candidatesTokenCount"` + } `json:"usageMetadata"` +} + +// Stream 实现 Adaptor：Gemini SSE 数据帧转 IR Delta。 +func (g *GeminiAdaptor) Stream(ctx context.Context, req *ir.ChatRequest, out chan<- ir.Delta) error { + body := gemRequest{} + for _, m := range req.Messages { + switch m.Role { + case "system": + if body.SystemInstruction == nil { + body.SystemInstruction = &gemContent{} + } + body.SystemInstruction.Parts = append(body.SystemInstruction.Parts, gemPart{Text: m.Content}) + case "assistant": + body.Contents = append(body.Contents, gemContent{Role: "model", Parts: []gemPart{{Text: m.Content}}}) + default: // user / tool 结果按 user 侧回传（工具循环随百工 M5 完整化） + body.Contents = append(body.Contents, gemContent{Role: "user", Parts: []gemPart{{Text: m.Content}}}) + } + } + if req.Temperature != nil || req.MaxTokens != nil { + body.Generation = &gemGeneration{Temperature: req.Temperature, MaxTokens: req.MaxTokens} + } + + payload, err := json.Marshal(body) + if err != nil { + return err + } + url := strings.TrimRight(g.BaseURL, "/") + + "/v1beta/models/" + g.UpstreamModel + ":streamGenerateContent?alt=sse" + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload)) + if err != nil { + return err + } + httpReq.Header.Set("content-type", "application/json") + httpReq.Header.Set("x-goog-api-key", g.APIKey) + + client := g.Client + if client == nil { + client = &http.Client{Timeout: 120 * time.Second} + } + resp, err := client.Do(httpReq) + if err != nil { + return &UpstreamError{StatusCode: 0, Message: err.Error()} + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode >= 400 { + buf := new(bytes.Buffer) + _, _ = buf.ReadFrom(resp.Body) + return &UpstreamError{StatusCode: resp.StatusCode, Message: buf.String()} + } + + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data:") { + continue + } + data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if data == "" { + continue + } + var chunk gemChunk + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + continue + } + if len(chunk.Candidates) > 0 { + c := chunk.Candidates[0] + for _, part := range c.Content.Parts { + if part.Text != "" { + out <- ir.Delta{Delta: part.Text} + } + } + if c.FinishReason != nil { + finish := "stop" + if strings.EqualFold(*c.FinishReason, "SAFETY") || strings.EqualFold(*c.FinishReason, "RECITATION") { + finish = "content_filter" + } + var up, uc int + if chunk.UsageMetadata != nil { + up, uc = chunk.UsageMetadata.PromptTokenCount, chunk.UsageMetadata.CandidatesTokenCount + } + out <- ir.Delta{FinishReason: finish, UsagePromptTokens: up, UsageCompletionTokens: uc} + } + } + } + return scanner.Err() +} diff --git a/shenshi/main.go b/shenshi/main.go index 7f91e94..21ba7d5 100644 --- a/shenshi/main.go +++ b/shenshi/main.go @@ -1,14 +1,19 @@ -// 神识（Shenshi）入口：魂镜模型网关，rewrite/1.0 M0 骨架。 -// 数据面 POST /v1/chat/completions（IR 入站，SSE 出站）；健康面 /healthz /readyz。 +// 神识（Shenshi）入口：魂镜模型网关（rewrite/1.0）。 +// 数据面 POST /v1/chat/completions（IR 入站，SSE 出站）；健康面 /healthz /readyz； +// 计量面 /me/usage、对账抽样 /admin/usage/sample；BYOK 面 /me/byok-keys。 package main import ( + "fmt" "log" "net/http" "os" + "hunjing/shenshi/internal/byok" + "hunjing/shenshi/internal/quota" "hunjing/shenshi/internal/route" "hunjing/shenshi/internal/server" + "hunjing/shenshi/internal/store" ) func envOr(key, fallback string) string { @@ -19,13 +24,38 @@ func envOr(key, fallback string) string { } func main() { - port := envOr("SHENSHI_PORT", "8080") - srv := server.New( - server.Config{ServiceToken: os.Getenv("SHENSHI_SERVICE_TOKEN")}, - route.New(), - ) - addr := ":" + port - log.Printf("神识就绪：http://0.0.0.0%s（stub 端点，等待 M4 接入真实 Provider）", addr) + st, err := store.Open(os.Getenv("SHENSHI_DATABASE_URL")) + if err != nil { + log.Fatalf("神识存储初始化失败: %v", err) + } + + cfg := server.Config{ + ServiceToken: os.Getenv("SHENSHI_SERVICE_TOKEN"), + Store: st, + } + + // BYOK：主密钥未配置时整体禁用（fail-closed） + if master := os.Getenv("SHENSHI_BYOK_MASTER_KEY"); master != "" { + vault, err := byok.NewVault(master, st) + if err != nil { + log.Fatalf("BYOK 主密钥配置无效: %v", err) + } + cfg.Vault = vault + log.Printf("BYOK 密钥库就绪") + } + + // 配额：用户级软配额（token 口径；0 = 不限制） + var budget int64 + if raw := os.Getenv("SHENSHI_USER_BUDGET_TOKENS"); raw != "" { + if _, err := fmt.Sscanf(raw, "%d", &budget); err != nil { + log.Fatalf("SHENSHI_USER_BUDGET_TOKENS 须为整数: %v", err) + } + cfg.Quota = quota.New(st, budget) + } + + srv := server.New(cfg, route.New()) + addr := ":" + envOr("SHENSHI_PORT", "8080") + log.Printf("神识就绪：http://0.0.0.0%s（端点目录 SHENSHI_ENDPOINTS，未配置时回声桩兜底）", addr) if err := http.ListenAndServe(addr, srv); err != nil { log.Fatal(err) }