🦀 技術調査 - agentgateway / kgateway による AI Gateway 拡張

kind create cluster --name ai-gateway-lab
kubectl cluster-info --context kind-ai-gateway-lab

kubectl apply -f https://github.com/kubernetes-sigs/gateway-api/releases/download/v1.2.0/standard-install.yaml
kubectl wait --for=condition=Established crd/gateways.gateway.networking.k8s.io --timeout=60s

helm upgrade -i kgateway-crds oci://cr.kgateway.dev/kgateway-dev/charts/kgateway-crds \
  --namespace kgateway-system --create-namespace --version v2.2.1
helm upgrade -i kgateway oci://cr.kgateway.dev/kgateway-dev/charts/kgateway \
  --namespace kgateway-system --version v2.2.1 \
  --set agentgateway.enabled=true
kubectl -n kgateway-system wait --for=condition=Ready pod -l app=kgateway --timeout=120s

helm upgrade -i agentgateway-crds oci://cr.agentgateway.dev/charts/agentgateway-crds \
  --namespace agentgateway-system --create-namespace --version v1.1.0
helm upgrade -i agentgateway oci://cr.agentgateway.dev/charts/agentgateway \
  --namespace agentgateway-system --version v1.1.0
kubectl -n agentgateway-system wait --for=condition=Ready pod -l app=agentgateway --timeout=120s

[package]
name = "my-rustformation"
version = "0.1.0"
edition = "2021"

[lib]
name = "my_rustformation"
crate-type = ["cdylib"]

[dependencies]
envoy-proxy-dynamic-modules-rust-sdk = "0.1"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
minijinja = { version = "2.12.0", features = ["loader"] }
anyhow = "1.0"

use envoy_proxy_dynamic_modules_rust_sdk::*;

pub struct MyTransformationFilter {
    request_id: Option<String>,
}

impl HttpFilter for MyTransformationFilter {
    fn on_request_headers(
        &mut self,
        _num_headers: usize,
        _end_of_stream: bool,
        ops: &dyn RequestHeadersOps,
    ) -> Result<FilterHeadersStatus> {
        ops.set_header("x-custom-team-id", "team-alpha")?;
        if let Ok(Some(id)) = ops.get_header("x-request-id") {
            self.request_id = Some(String::from_utf8_lossy(id).to_string());
        }
        Ok(FilterHeadersStatus::Continue)
    }

    fn on_request_body(
        &mut self,
        _data_size: usize,
        end_of_stream: bool,
        ops: &dyn RequestBodyOps,
    ) -> Result<FilterDataStatus> {
        if !end_of_stream {
            return Ok(FilterDataStatus::StopIterationAndBuffer);
        }
        let body = ops.get_body()?;
        let mut json: serde_json::Value = serde_json::from_slice(body)?;
        if let Some(messages) = json.get_mut("messages") {
            // PII マスキング
        }
        let transformed = serde_json::to_vec(&json)?;
        ops.replace_body(&transformed)?;
        Ok(FilterDataStatus::Continue)
    }

    fn on_response_headers(
        &mut self,
        _num_headers: usize,
        _end_of_stream: bool,
        ops: &dyn ResponseHeadersOps,
    ) -> Result<FilterHeadersStatus> {
        ops.add_header("x-gateway-version", "agentgateway-1.1.0")?;
        Ok(FilterHeadersStatus::Continue)
    }

    fn on_response_body(
        &mut self,
        _data_size: usize,
        end_of_stream: bool,
        ops: &dyn ResponseBodyOps,
    ) -> Result<FilterDataStatus> {
        if !end_of_stream {
            return Ok(FilterDataStatus::StopIterationAndBuffer);
        }
        let body = ops.get_body()?;
        if let Ok(json) = serde_json::from_slice::<serde_json::Value>(body) {
            if let Some(tokens) = json.pointer("/usage/total_tokens") {
                eprintln!("[rustformation] total_tokens={}", tokens);
            }
        }
        Ok(FilterDataStatus::Continue)
    }

    fn on_exchange_complete(&mut self, _ops: &dyn ExchangeCompleteOps) -> Result<()> {
        Ok(())
    }
}

#[no_mangle]
pub extern "C" fn envoy_dynamic_module_on_program_init() -> *const std::os::raw::c_char {
    std::ptr::null()
}

FROM rust:1.85 AS builder
RUN apt-get update && apt-get install -y clang lld && rm -rf /var/lib/apt/lists/*
WORKDIR /build
COPY Cargo.toml Cargo.lock ./
COPY src ./src
RUN cargo build --release --target x86_64-unknown-linux-gnu

FROM envoyproxy/envoy:v1.36.4
COPY --from=builder \
  /build/target/x86_64-unknown-linux-gnu/release/libmy_rustformation.so \
  /usr/local/lib/
ENV ENVOY_DYNAMIC_MODULES_SEARCH_PATH=/usr/local/lib

apiVersion: gateway.networking.k8s.io/v1
kind: Gateway
metadata:
  name: ai-gateway
  namespace: default
spec:
  gatewayClassName: kgateway
  listeners:
    - name: http
      protocol: HTTP
      port: 80

apiVersion: agentgateway.dev/v1alpha1
kind: AgentgatewayBackend
metadata:
  name: openai-backend
  namespace: agentgateway-system
spec:
  ai:
    provider:
      openai:
        model: gpt-4
      host: api.openai.com
      port: 443
      path: "/v1/chat/completions"

apiVersion: gateway.networking.k8s.io/v1
kind: HTTPRoute
metadata:
  name: llm-route
  namespace: default
spec:
  parentRefs:
    - name: ai-gateway
      namespace: default
  rules:
    - matches:
        - path:
            type: PathPrefix
            value: /v1/chat/completions
      backendRefs:
        - group: agentgateway.dev
          kind: AgentgatewayBackend
          name: openai-backend

apiVersion: gateway.kgateway.dev/v1alpha1
kind: TrafficPolicy
metadata:
  name: rust-transformation-policy
  namespace: default
spec:
  targetRefs:
    - group: gateway.networking.k8s.io
      kind: HTTPRoute
      name: llm-route
  # dynamicModules フィールドの正確な構文は kgateway リリースノートを参照

GATEWAY_URL=$(kubectl get gateway ai-gateway -o jsonpath='{.status.addresses[0].value}')

curl -X POST "http://${GATEWAY_URL}/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -H "x-request-id: test-001" \
  -d '{
    "model": "gpt-4o",
    "messages": [{"role": "user", "content": "Hello, world!"}],
    "stream": false
  }'

curl -X POST "http://${GATEWAY_URL}/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -H "Accept: text/event-stream" \
  -d '{
    "model": "gpt-4o",
    "messages": [{"role": "user", "content": "Explain Rust in one sentence."}],
    "stream": true
  }'

kubectl get pods -n agentgateway -l app=agentgateway
kubectl rollout restart deployment/agentgateway -n agentgateway
kubectl rollout status deployment/agentgateway -n agentgateway
kubectl scale deployment/agentgateway --replicas=0 -n agentgateway

kubectl port-forward pod/<pod-name> 9095:9095 -n agentgateway
curl http://localhost:9095/config_dump | jq .
curl http://localhost:9095/stats | grep dynamic_module
curl http://localhost:9095/ready

livenessProbe:
  httpGet:
    path: /ready
    port: 9095
  initialDelaySeconds: 10
  periodSeconds: 10
readinessProbe:
  httpGet:
    path: /ready
    port: 9095
  initialDelaySeconds: 5
  periodSeconds: 5

access_log:
  - name: envoy.access_loggers.file
    typed_config:
      "@type": type.googleapis.com/envoy.extensions.access_loggers.file.v3.FileAccessLog
      path: /dev/stdout
      log_format:
        json_format:
          timestamp: "%START_TIME%"
          method: "%REQ(:METHOD)%"
          response_code: "%RESPONSE_CODE%"
          duration_ms: "%DURATION%"
          x_request_id: "%REQ(X-REQUEST-ID)%"

telemetry:
  tracing:
    enabled: true
    exporter:
      otlp:
        endpoint: "http://otel-collector:4317"
  metrics:
    enabled: true
    exporter:
      prometheus:
        port: 9090

spec:
  strategy:
    type: RollingUpdate
    rollingUpdate:
      maxSurge: 1
      maxUnavailable: 0
  template:
    spec:
      terminationGracePeriodSeconds: 60

kubectl scale deployment/agentgateway --replicas=5 -n agentgateway
kubectl autoscale deployment/agentgateway \
  --cpu-percent=70 --min=2 --max=10 -n agentgateway

security:
  auditLog:
    enabled: true
    format: json
    tamperEvident: true
    destination: stdout
    fields:
      - request_id
      - user_id
      - model
      - token_usage
      - prompt_hash
      - response_latency_ms

location /v1/chat/completions {
    proxy_pass http://agentgateway;
    proxy_buffering off;
    proxy_cache off;
    proxy_read_timeout 300s;
    tcp_nodelay on;
}

use std::panic::{catch_unwind, AssertUnwindSafe};

fn on_request_headers(...) -> Result<FilterHeadersStatus> {
    let result = catch_unwind(AssertUnwindSafe(|| {
        self.do_request_headers(num_headers, end_of_stream, ops)
    }));
    match result {
        Ok(Ok(status)) => Ok(status),
        Ok(Err(e)) => {
            eprintln!("[rustformation] ERROR err={}", e);
            Ok(FilterHeadersStatus::Continue)
        }
        Err(_) => {
            eprintln!("[rustformation] PANIC — fail-closed");
            Ok(FilterHeadersStatus::StopIteration)
        }
    }
}

jobs:
  version-matrix:
    strategy:
      matrix:
        envoy_version: ["v1.36.4"]
        rust_version: ["1.85"]
    steps:
      - name: Build .so
        run: |
          docker build \
            --build-arg ENVOY_VERSION=${{ matrix.envoy_version }} \
            --build-arg RUST_VERSION=${{ matrix.rust_version }} \
            -t agentgateway-ext:ci .

securityContext:
  runAsNonRoot: true
  runAsUser: 1000
  readOnlyRootFilesystem: true
  capabilities:
    drop:
      - ALL
    add:
      - NET_BIND_SERVICE
volumeMounts:
  - name: rust-module
    mountPath: /usr/local/lib
    readOnly: true