async def get_dynamic_limit(request: Request, db: Session = Depends(get_session)):
    deployment_id = request.path_params.get("id")
    # 1. 优先查 Redis 缓存的配置 (减少 DB 压力)
    # 2. Redis 没有则查 DB
    deployment = ... 
    return deployment.rate_limit or "60/minute"

# tests/agent_server/test_deployment.py
import pytest
from httpx import AsyncClient

@pytest.mark.asyncio
async def test_publish_flow_versioning(client: AsyncClient, json_flow_data):
    # 1. 第一次发布
    resp1 = await client.post("/api/v1/deployments/publish", json={
        "flow_id": "flow-123",
        "endpoint_slug": "test-agent",
        "name": "Test Agent"
    })
    assert resp1.status_code == 200
    assert resp1.json()["version"] == 1

    # 2. 第二次发布
    resp2 = await client.post("/api/v1/deployments/publish", json={
        "flow_id": "flow-123",
        "endpoint_slug": "test-agent",
        "name": "Test Agent Updated"
    })
    assert resp2.status_code == 200
    assert resp2.json()["version"] == 2
    assert resp2.json()["id"] == resp1.json()["id"] # ID 相同，版本不同

# tests/agent_server/test_rate_limit.py
import pytest
import asyncio
from httpx import AsyncClient

@pytest.mark.asyncio
async def test_rate_limiting_enforcement(client: AsyncClient, active_deployment):
    # 假设 active_deployment 限流设置为 "2/second"
    slug = active_deployment.endpoint_slug
    
    # 并发发送 3 个请求
    tasks = [client.post(f"/api/v1/agent/{slug}", json={"input": "hi"}) for _ in range(3)]
    responses = await asyncio.gather(*tasks)
    
    status_codes = [r.status_code for r in responses]
    # 应该有两个 200，一个 429
    assert 429 in status_codes
    assert status_codes.count(200) <= 2

# tests/agent_server/test_execution.py
from unittest.mock import patch

@pytest.mark.asyncio
async def test_execution_cache(agent_service, deployment_obj):
    # Mock 加载构建过程
    with patch.object(agent_service, 'load_and_build_graph', wraps=agent_service.load_and_build_graph) as mock_build:
        # 第一次执行
        await agent_service.execute_agent(deployment_obj, {"input": "1"})
        # 第二次执行
        await agent_service.execute_agent(deployment_obj, {"input": "2"})
        
        # 断言构建只发生了一次
        assert mock_build.call_count == 1

from datetime import datetime
from typing import Optional, Dict, Any, List
from uuid import UUID, uuid4
from sqlmodel import Field, SQLModel, JSON, Column, Relationship

# 部署主表：控制对外接口、限流和当前激活版本
class Deployment(SQLModel, table=True):
    id: UUID = Field(default_factory=uuid4, primary_key=True)
    flow_id: UUID = Field(index=True)
    name: str
    endpoint_slug: str = Field(unique=True, index=True)
    description: Optional[str] = None
    rate_limit: str = Field(default="60/minute")
    is_active: bool = Field(default=True)
    active_version: int = Field(default=1) # 当前对外服务的版本号
    
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    
    versions: List["DeploymentVersion"] = Relationship(back_populates="deployment")

# 版本快照表：存储每次发布的静态数据
class DeploymentVersion(SQLModel, table=True):
    id: UUID = Field(default_factory=uuid4, primary_key=True)
    deployment_id: UUID = Field(foreign_key="deployment.id")
    version: int
    graph_snapshot: Dict[str, Any] = Field(default={}, sa_column=Column(JSON)) # 冻结的图数据
    production_config: Dict[str, Any] = Field(default={}, sa_column=Column(JSON)) # 覆盖配置
    created_at: datetime = Field(default_factory=datetime.utcnow)
    
    deployment: Optional[Deployment] = Relationship(back_populates="versions")

# Pydantic Models for API
class DeploymentCreate(SQLModel):
    flow_id: UUID
    name: str
    endpoint_slug: str
    production_config: Optional[Dict[str, Any]] = None

class DeploymentRead(SQLModel):
    id: UUID
    name: str
    endpoint_slug: str
    active_version: int
    rate_limit: str
    is_active: bool
    updated_at: datetime

import logging
import os
import json
import asyncio
from typing import Any, Dict, Optional, Tuple, AsyncGenerator
from uuid import UUID

from async_lru import alru_cache
from langchain_core.runnables import Runnable, RunnableConfig
from langfuse.callback import CallbackHandler
from sqlmodel import Session, select

from langflow.graph.graph.base import Graph
from langflow.services.database.models.agent_server import Deployment, DeploymentVersion
from langflow.services.get_service import get_session

logger = logging.getLogger("agent_server")

class AgentServerService:
    _cache = {} # 内存缓存 (DeploymentID_Version -> Runnable)

    async def get_runnable(self, deployment: Deployment, db: Session) -> Runnable:
        """获取可执行对象 (带缓存)"""
        # 查找当前激活的版本快照
        version_record = db.exec(
            select(DeploymentVersion)
            .where(DeploymentVersion.deployment_id == deployment.id, 
                   DeploymentVersion.version == deployment.active_version)
        ).first()
        
        if not version_record:
            raise ValueError(f"Version {deployment.active_version} not found for deployment {deployment.name}")

        cache_key = f"{deployment.id}_v{version_record.version}"
        
        if cache_key in self._cache:
            return self._cache[cache_key]

        # Cache Miss: 构建图
        logger.info(f"Building graph for {cache_key}")
        graph_data = version_record.graph_snapshot
        prod_config = version_record.production_config

        # 1. 从 Snapshot 恢复 Graph
        graph = Graph.from_payload(graph_data)
        
        # 2. 应用生产环境配置 (Tweaks)
        if prod_config:
            for node in graph.vertices:
                if node.id in prod_config:
                    for key, val in prod_config[node.id].items():
                        node.params[key] = val
        
        # 3. 编译
        # 注意：这里适配 Langflow 的 build 逻辑
        built_artifacts = await graph.build() if asyncio.iscoroutinefunction(graph.build) else graph.build()
        
        # 4. 提取 Runnable (兼容 LangChain Chain 和 LangGraph CompiledGraph)
        runnable = built_artifacts
        if isinstance(built_artifacts, list):
            # 尝试寻找主要入口
            runnable = built_artifacts[0] 
        
        # 写入缓存
        self._cache[cache_key] = runnable
        return runnable

    async def execute_stream(
        self, 
        deployment: Deployment, 
        inputs: Dict[str, Any], 
        session_id: Optional[str],
        db: Session
    ) -> AsyncGenerator[str, None]:
        """流式执行 (SSE)"""
        try:
            runnable = await self.get_runnable(deployment, db)
            
            # 准备 Langfuse
            callbacks = []
            lf_handler = self._create_langfuse(
                trace_name=f"agent-{deployment.name}",
                session_id=session_id,
                tags=[deployment.endpoint_slug, f"v{deployment.active_version}"]
            )
            if lf_handler:
                callbacks.append(lf_handler)

            config = RunnableConfig(
                callbacks=callbacks,
                configurable={"session_id": session_id}
            )

            # 判断是否支持 astream_events (LangGraph v0.1+)
            if hasattr(runnable, "astream_events"):
                async for event in runnable.astream_events(inputs, config=config, version="v2"):
                    event_type = event["event"]
                    data = event["data"]

                    if event_type == "on_chat_model_stream":
                        chunk = data.get("chunk")
                        if chunk and hasattr(chunk, "content") and chunk.content:
                            yield self._sse_pack("token", {"content": chunk.content})
                    
                    elif event_type == "on_tool_start":
                        yield self._sse_pack("tool_start", {"tool": event["name"]})
                    
                    elif event_type == "on_tool_end":
                        yield self._sse_pack("tool_end", {"output": str(data.get("output"))})
            
            else:
                # 普通 Chain 的流式回退
                async for chunk in runnable.astream(inputs, config=config):
                    content = chunk.content if hasattr(chunk, "content") else str(chunk)
                    yield self._sse_pack("token", {"content": content})

            yield self._sse_pack("end", "[DONE]")

        except Exception as e:
            logger.error(f"Execution error: {e}", exc_info=True)
            yield self._sse_pack("error", str(e))

    def _sse_pack(self, event: str, data: Any) -> str:
        return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"

    def _create_langfuse(self, **kwargs):
        pk = os.getenv("LANGFUSE_PUBLIC_KEY")
        sk = os.getenv("LANGFUSE_SECRET_KEY")
        if pk and sk:
            return CallbackHandler(public_key=pk, secret_key=sk, host=os.getenv("LANGFUSE_HOST"), **kwargs)
        return None

agent_service = AgentServerService()

from typing import List
from fastapi import APIRouter, Depends, HTTPException, Body
from sqlmodel import Session, select
from langflow.services.get_service import get_session
from langflow.services.database.models.flow import Flow
from langflow.services.database.models.agent_server import Deployment, DeploymentVersion, DeploymentCreate, DeploymentRead

router = APIRouter(prefix="/deployments", tags=["Agent Management"])

@router.post("/publish", response_model=DeploymentRead)
def publish_flow(data: DeploymentCreate, session: Session = Depends(get_session)):
    """发布新版本"""
    # 1. 获取 Flow 数据
    flow = session.get(Flow, data.flow_id)
    if not flow:
        raise HTTPException(404, "Flow not found")
    
    # 2. 检查或创建 Deployment
    deployment = session.exec(select(Deployment).where(Deployment.endpoint_slug == data.endpoint_slug)).first()
    
    if not deployment:
        deployment = Deployment(
            flow_id=data.flow_id,
            name=data.name,
            endpoint_slug=data.endpoint_slug,
            active_version=1
        )
        session.add(deployment)
        session.commit()
        session.refresh(deployment)
        new_version_num = 1
    else:
        # 版本自增
        last_version = session.exec(
            select(DeploymentVersion)
            .where(DeploymentVersion.deployment_id == deployment.id)
            .order_by(DeploymentVersion.version.desc())
        ).first()
        new_version_num = (last_version.version + 1) if last_version else 1
        
        # 更新 Deployment 信息
        deployment.active_version = new_version_num
        deployment.name = data.name
        session.add(deployment)

    # 3. 创建版本快照
    version = DeploymentVersion(
        deployment_id=deployment.id,
        version=new_version_num,
        graph_snapshot=flow.data.copy(), # 深度复制
        production_config=data.production_config or {}
    )
    session.add(version)
    session.commit()
    return deployment

@router.patch("/{id}/config")
def update_config(id: str, rate_limit: str = Body(..., embed=True), is_active: bool = Body(..., embed=True), session: Session = Depends(get_session)):
    """更新限流配置或停用服务"""
    dep = session.get(Deployment, id)
    if not dep:
        raise HTTPException(404)
    dep.rate_limit = rate_limit
    dep.is_active = is_active
    session.add(dep)
    session.commit()
    return {"status": "updated"}

@router.post("/{id}/rollback")
def rollback_version(id: str, version: int = Body(..., embed=True), session: Session = Depends(get_session)):
    """版本回退"""
    dep = session.get(Deployment, id)
    ver = session.exec(select(DeploymentVersion).where(DeploymentVersion.deployment_id == id, DeploymentVersion.version == version)).first()
    if not dep or not ver:
        raise HTTPException(404, "Deployment or Version not found")
    
    dep.active_version = version
    session.add(dep)
    session.commit()
    return {"status": "rolled_back", "current_version": version}

@router.get("/flow/{flow_id}", response_model=List[DeploymentRead])
def get_deployments_by_flow(flow_id: str, session: Session = Depends(get_session)):
    return session.exec(select(Deployment).where(Deployment.flow_id == flow_id)).all()

from fastapi import APIRouter, Depends, HTTPException, Request, Body
from fastapi.responses import StreamingResponse
from fastapi_limiter.depends import RateLimiter
from sqlmodel import Session, select
from langflow.services.get_service import get_session
from langflow.services.database.models.agent_server import Deployment
from langflow.services.agent_server.service import agent_service

router = APIRouter(prefix="/agent", tags=["Agent Runtime"])

async def get_dynamic_rate_limit(request: Request, session: Session = Depends(get_session)):
    """从数据库动态获取限流规则"""
    slug = request.path_params.get("slug")
    dep = session.exec(select(Deployment).where(Deployment.endpoint_slug == slug)).first()
    return dep.rate_limit if dep else "5/minute" # 默认严格限流

async def rate_limit_key(request: Request):
    return f"{request.path_params.get('slug')}:{request.client.host}"

@router.post("/{slug}/stream", dependencies=[Depends(RateLimiter(key_func=rate_limit_key, limit_func=get_dynamic_rate_limit))])
async def stream_agent(
    slug: str,
    inputs: dict = Body(...),
    session_id: str = Body(None),
    db: Session = Depends(get_session)
):
    """对外流式接口"""
    deployment = db.exec(select(Deployment).where(Deployment.endpoint_slug == slug, Deployment.is_active == True)).first()
    if not deployment:
        raise HTTPException(404, "Agent not found or inactive")

    return StreamingResponse(
        agent_service.execute_stream(deployment, inputs, session_id, db),
        media_type="text/event-stream"
    )

# ... existing imports
import redis.asyncio as redis
from fastapi_limiter import FastAPILimiter
from langflow.api.v1.agent_mgmt import router as mgmt_router
from langflow.api.v1.agent_runtime import router as runtime_router

# ... inside create_app
@app.on_event("startup")
async def startup():
    # Redis for Rate Limiting
    redis_url = os.getenv("REDIS_URL", "redis://localhost:6379")
    r = redis.from_url(redis_url, encoding="utf-8", decode_responses=True)
    await FastAPILimiter.init(r)

app.include_router(mgmt_router)
app.include_router(runtime_router)

import pytest
from httpx import AsyncClient

# 假设已经配置好了 conftest.py 提供 client fixture

@pytest.mark.asyncio
async def test_end_to_end_flow(client: AsyncClient, active_user):
    # 1. 发布流程
    # 假设我们有一个存在的 flow_id
    flow_id = "existing-flow-uuid" 
    
    resp = await client.post("/api/v1/deployments/publish", json={
        "flow_id": flow_id,
        "name": "Integration Test Agent",
        "endpoint_slug": "test-integration"
    })
    assert resp.status_code == 200
    dep_data = resp.json()
    assert dep_data["version"] == 1
    
    # 2. 验证限流 (初始)
    # 调用 Runtime 接口
    slug = "test-integration"
    resp = await client.post(f"/api/v1/agent/{slug}/stream", json={"input_value": "hi"})
    # 应该是 200 (StreamingResponse start)
    assert resp.status_code == 200 

    # 3. 修改限流配置
    dep_id = dep_data["id"]
    await client.patch(f"/api/v1/deployments/{dep_id}/config", json={
        "rate_limit": "0/minute", # 禁止访问
        "is_active": True
    })
    
    # 4. 再次调用应该被拦截
    resp = await client.post(f"/api/v1/agent/{slug}/stream", json={"input_value": "hi"})
    assert resp.status_code == 429 # Too Many Requests

    # 5. 版本回退测试
    # 再次发布产生 v2
    await client.post("/api/v1/deployments/publish", json={
        "flow_id": flow_id, "name": "v2", "endpoint_slug": "test-integration"
    })
    
    # 执行回退到 v1
    rollback_resp = await client.post(f"/api/v1/deployments/{dep_id}/rollback", json={"version": 1})
    assert rollback_resp.json()["current_version"] == 1