diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 8a13e914..931cce20 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -474,6 +474,25 @@ path. Enabling the route is an explicit configuration change; it never selects or changes a provider or model. `quota-axi` remains an observation-only tool and has no routing, login or mutation authority. +An approved Responses client that cannot supply command-backed authentication +uses the AgentOS `ai-gateway-workload-proxy` as a loopback-only sidecar; the +exact Hermes contract and Pod wiring live in +[`docs/integrations/hermes-ai-gateway.md`](./docs/integrations/hermes-ai-gateway.md). +The sidecar rereads projected workload identity for each request, removes the +client's placeholder authorization and all caller-supplied AgentOS authority +headers, optionally injects a validated Assignment ID from its own trusted +configuration, and forwards only the two reviewed Responses paths once to +Agentgateway. It preserves the upstream status, end-to-end headers after +required hop-by-hop and stale `Content-Encoding`/`Content-Length` metadata +removal, and stream without following redirects. W3C trace correlation passes +through unchanged. Its +process-only liveness and local projected-token readiness never contact an +upstream. It owns no credential, +selection, retry, fallback, prompt, response, or deployment authority. The +client keeps exact model selection and disables request replay and provider +fallback for this route; its workload owner composes the sidecar and projected +token into that client's reviewed native manifest. + This is not a universal AgentOS proxy. Git, PostgreSQL, Kubernetes, Herdr, registries and other provider tools continue through their native interfaces. Adding another mediated protocol requires a separate review of its authority, diff --git a/Dockerfile b/Dockerfile index e26941bb..5cf716d4 100644 --- a/Dockerfile +++ b/Dockerfile @@ -225,6 +225,7 @@ RUN chmod 0644 \ /opt/agentos/services/a2a/src/main.ts \ /opt/agentos/services/agentgateway/src/readiness-main.ts \ /opt/agentos/services/ai-gateway/src/main.ts \ + /opt/agentos/services/ai-gateway/src/workload-client-proxy-main.ts \ /opt/agentos/services/egress-authz/src/main.ts \ /opt/agentos/services/github-broker/src/main.ts \ /opt/agentos/services/openfga/src/bootstrap.ts \ @@ -247,6 +248,9 @@ RUN chmod 0644 \ && ln -s \ /opt/agentos/services/ai-gateway/src/main.ts \ /usr/local/bin/ai-gateway \ + && ln -s \ + /opt/agentos/services/ai-gateway/src/workload-client-proxy-main.ts \ + /usr/local/bin/ai-gateway-workload-proxy \ && ln -s \ /opt/agentos/services/egress-authz/src/main.ts \ /usr/local/bin/agentos-egress-authz \ diff --git a/docs/integrations/hermes-ai-gateway.config.yaml b/docs/integrations/hermes-ai-gateway.config.yaml new file mode 100644 index 00000000..3b0c9941 --- /dev/null +++ b/docs/integrations/hermes-ai-gateway.config.yaml @@ -0,0 +1,20 @@ +# Hermes Agent 0.20.0 (v2026.8.3) deterministic integration fixture. +model: + provider: custom:agentos-gateway + default: gpt-exact-fixture + +providers: + agentos-gateway: + name: AgentOS Gateway + api: http://127.0.0.1:8790/v1 + api_key: agentos-workload-identity-placeholder + transport: codex_responses + discover_models: false + models: + - gpt-exact-fixture + +agent: + api_max_retries: 0 + +fallback_providers: [] +fallback_model: "" diff --git a/docs/integrations/hermes-ai-gateway.md b/docs/integrations/hermes-ai-gateway.md new file mode 100644 index 00000000..1e39ff31 --- /dev/null +++ b/docs/integrations/hermes-ai-gateway.md @@ -0,0 +1,193 @@ +# Hermes Agent through the Fleet AI Gateway + +This contract covers Hermes Agent `0.20.0` from release tag `v2026.8.3` and +other approved clients that speak the OpenAI Responses API but cannot reread a +Kubernetes projected token for every request. Direct per-agent OAuth remains +the recovery path. + +> Deployment gate: this Hermes revision is not eligible for this route. Its +> `codex_responses` runtime hardcodes one retry around `responses.create`, so a +> transport failure after partial output can replay the request; +> `HERMES_STREAM_RETRIES=0` does not change that behavior. Do not deploy the +> fixture or Pod wiring below until a pinned Hermes revision with a verified +> zero-retry Responses path is available. + +## Client contract + +The client keeps authority over the exact model and reasoning configuration. It +sends only `POST /v1/responses` or `POST /v1/responses/compact` to a loopback +`ai-gateway-workload-proxy` sidecar. The sidecar: + +- accepts no provider credential and listens on loopback only; +- rereads the kubelet-rotated `agentos-egress-authz` ServiceAccount token for + every request; +- removes caller-supplied authorization, assignment, decision, grant, and every + other `x-agentos-*` header, then supplies only projected identity and an + optional validated assignment ID from sidecar-owned configuration to + `agentgateway-openai`; W3C `traceparent` and `tracestate` correlation remain + unchanged; +- forwards the request once, without inspecting or storing its body; and +- returns the actual upstream status, end-to-end headers after required + hop-by-hop and stale `Content-Encoding`/`Content-Length` metadata removal, + and stream without retrying, following redirects, or selecting another model + or account. A transport failure returns the adapter's stable `502`; it never + becomes apparent success. + +Agentgateway and `agentos-egress-authz`, not the sidecar, authenticate the Pod +and authorize its registered Agent or Assignment. Do not expose the loopback +listener as a Service. Do not copy an AI Gateway client Secret, OAuth file, or +operator token into the client namespace. + +`GET /livez` is process-only. `GET /readyz` rereads the projected token and +returns ready only when it is present, at most 16 KiB, strict UTF-8, already +trimmed, and JWT-like. Neither endpoint contacts Agentgateway or a provider. + +## Hermes 0.20.0 configuration + +Hermes 0.20 supports named providers with `transport: codex_responses`, but a +configured API key is a process-lifetime value. Point the named provider at the +loopback adapter and use a non-secret placeholder; never put the projected token +in `config.yaml` or `.env`. + +```yaml +model: + provider: custom:agentos-gateway + default: + +providers: + agentos-gateway: + name: AgentOS Gateway + api: http://127.0.0.1:8790/v1 + api_key: agentos-workload-identity-placeholder + transport: codex_responses + discover_models: false + models: + - + +agent: + api_max_retries: 0 + +fallback_providers: [] +fallback_model: "" +``` + +These fields are load-bearing for Hermes 0.20's named-provider resolver. The +input must select `provider: custom:agentos-gateway`; the matching provider entry's +`transport: codex_responses` retains the Responses transport, and its `api` +points that transport at the loopback `/v1` base. Hermes normalizes this named +input at runtime to the named custom provider with `api_mode: codex_responses` +and `base_url: http://127.0.0.1:8790/v1`. Putting bare `provider: custom` +directly in the input falls back to Chat Completions. The checked fixture +[`hermes-ai-gateway.config.yaml`](./hermes-ai-gateway.config.yaml) is exercised +against the named input contract and the sidecar allowlist in the Gateway test +suite. + +`api_max_retries: 0` and the empty fallback configuration are part of the +contract: after Hermes sends a request, it must surface the real `401`, `403`, +`429`, timeout, or provider failure rather than replaying the turn through a +route that may acquire another account. Auxiliary requests are out of scope for +this blocked Hermes revision. Do not enable them or configure `provider: auto`, +`provider: custom`, or a fallback chain until Hermes provides a verified native +Responses path with no automatic retry, redirect, or fallback. + +## Pod wiring + +The owning Hermes deployment supplies this wiring. AgentOS owns the adapter +binary and the governed backend topology, but does not own or apply a foreign +Hermes StatefulSet. Add the selected-client label, projected identity volume, +and loopback sidecar to the reviewed workload manifest: + +The current Agentgateway NetworkPolicy admits client Pods from the `agentos` +namespace, or from an already approved namespace labeled +`agentos.akua.dev/managed-by: agentos-firstmate`, only when the Pod has an +`agentos.akua.dev/agent` label. Deploy Hermes in one of those approved +namespaces and preserve that reachability label. These labels only permit +network reachability; the projected ServiceAccount token and Agentgateway +authorizer provide authentication and authorization. + +```yaml +spec: + template: + metadata: + labels: + agentos.akua.dev/agent: "hermes" + agentos.akua.dev/agentgateway-client: "true" + spec: + automountServiceAccountToken: false + containers: + - name: hermes + - name: ai-gateway-workload-proxy + image: + command: ["ai-gateway-workload-proxy"] + env: + - name: AI_GATEWAY_URL + value: http://agentgateway-openai.agentos.svc.cluster.local:8788 + - name: AI_GATEWAY_IDLE_TIMEOUT_SECONDS + value: "255" + - name: AI_GATEWAY_GRACEFUL_SHUTDOWN_MILLIS + value: "20000" + - name: AGENTOS_EGRESS_TOKEN_FILE + value: /var/run/secrets/agentos-egress/token + # Optional: set only from the workload owner's trusted Assignment. + - name: AGENTOS_ASSIGNMENT_ID + value: + ports: + - name: workload-proxy + containerPort: 8790 + livenessProbe: + exec: + command: + - /usr/bin/curl + - --fail + - --silent + - --show-error + - http://127.0.0.1:8790/livez + readinessProbe: + exec: + command: + - /usr/bin/curl + - --fail + - --silent + - --show-error + - http://127.0.0.1:8790/readyz + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: ["ALL"] + readOnlyRootFilesystem: true + runAsNonRoot: true + volumeMounts: + - name: agentos-egress-identity + mountPath: /var/run/secrets/agentos-egress + readOnly: true + volumes: + - name: agentos-egress-identity + projected: + defaultMode: 0440 + sources: + - serviceAccountToken: + audience: agentos-egress-authz + expirationSeconds: 600 + path: token +``` + +Use the workload's existing explicit `fsGroup` so the sidecar can read the +`0440` projection. Only the sidecar mounts that identity volume; Hermes cannot +read it. Keep the workload's dedicated ServiceAccount, register that identity +and its approved access profile through the normal AgentOS access-plane +procedure, and preserve the existing NetworkPolicy. The label grants only network +reachability to Agentgateway; it is not authorization. + +The proxy's idle timeout defaults to the shown maximum of 255 seconds and its +graceful-shutdown timeout defaults to 20,000 milliseconds. The idle timeout is +an inactivity limit, not a total response limit; retain the 255-second value +when the approved provider stream can have gaps that long. The proxy uses +manual redirect handling, so a `3xx` response is returned to Hermes rather than +followed. + +## Operator workflow + +Use `$agentos-ai-gateway` for the approval, migration, rollout, verification, +rollback, and retirement workflow, including migration from the live v0.1.24 +shared-token topology. This page owns only the Hermes client contract and Pod +wiring. diff --git a/docs/security/workload-identity.md b/docs/security/workload-identity.md index ef42da14..0ef0172f 100644 --- a/docs/security/workload-identity.md +++ b/docs/security/workload-identity.md @@ -11,7 +11,7 @@ AgentOS workloads receive a second, explicit projected ServiceAccount token for - requested lifetime: 600 seconds; and - volume and mount: `agentos-egress-identity`, read-only with mode `0440`. -This is not the workload's Kubernetes API token. Persistent Second Mates retain their separately managed Kubernetes supervision identity. Crewmates keep `automountServiceAccountToken: false` and receive only the explicit audience-scoped projection. Only the long-running Agent container mounts this token; tool-install and home-preparation init containers do not, and domain admission rejects any Crewmate init-container mount. The kubelet rotates projected tokens, so clients must read the token file for each request rather than retaining its contents for the process lifetime. +This is not the workload's Kubernetes API token. Persistent Second Mates retain their separately managed Kubernetes supervision identity. Crewmates keep `automountServiceAccountToken: false` and receive only the explicit audience-scoped projection. Only the long-running provider client or a reviewed provider sidecar/adapter mounts this token; tool-install and home-preparation init containers do not, and domain admission rejects any Crewmate init-container mount. The kubelet rotates projected tokens, so the component that presents the identity must read the token file for each request rather than retaining its contents for the process lifetime. The authorization service built by #90 supplies the raw token only to `KubernetesTokenReviewer`. It requests exactly the dedicated audience and accepts a result only when `authenticated` is true and the same audience appears in `status.audiences`. Kubernetes documents that a TokenReview client setting audiences must check the returned intersection; accepting an empty result would fall back to API-server audience semantics and is forbidden here. See the Kubernetes [TokenReview API](https://kubernetes.io/docs/reference/kubernetes-api/definitions/token-review-v1-authentication/). diff --git a/package.json b/package.json index 0a25c22b..308d0651 100644 --- a/package.json +++ b/package.json @@ -19,7 +19,7 @@ "resilience:hard-gate": "bun packages/agentos/runtime/resilience-hard-gate-main.ts", "effect:tooling:test": "vitest run clis/pg-listen/tests/pg-listen.effect.test.ts release/kubernetes/tests/release.effect.test.ts", "effect:agentgateway:test": "vitest run services/agentgateway/tests/kubernetes.effect.test.ts services/agentgateway/tests/readiness.effect.test.ts", - "effect:test": "vitest run --no-file-parallelism mise.effect.test.ts tooling/effect-migration clis/github-app-token/tests/github-app-token.effect.test.ts packages/agentos/tests/*.effect.test.ts services/openfga/tests/*.effect.test.ts benchmarks/tests/contracts.effect.test.ts benchmarks/tests/pi-session-adapter.effect.test.ts benchmarks/tests/runner.effect.test.ts database/kubernetes/cloudnative-pg/tests/manifest.effect.test.ts database/tests/*.effect.test.ts packages/agentos/resources/crewmates/default/kubernetes/tests/manifest.effect.test.ts packages/agentos/resources/roles/firstmate/kubernetes/tests/manifests.effect.test.ts packages/agentos/resources/roles/secondmate/kubernetes/tests/*.effect.test.ts packages/agentos/runtime/kubernetes/tests/*.effect.test.ts packages/agentos/runtime/tests/kubernetes.effect.test.ts services/otel-collector/kubernetes/tests/*.effect.test.ts services/otel-collector/tests/*.effect.test.ts packages/agentos/src/shared/tests packages/agentos/src/access/tests packages/agentos/src/harness-control/tests packages/agentos/src/protocol/tests packages/agentos/src/resilience/tests packages/agentos/src/topology/tests packages/agentos/src/telemetry/tests packages/agentos/src/openai-server-compaction/tests/extension.effect.test.ts packages/agentos/src/openai-server-compaction/tests/messages.effect.test.ts packages/agentos/src/openai-server-compaction/tests/remote-runtime.effect.test.ts packages/agentos/src/openai-server-compaction/tests/schemas.effect.test.ts packages/agentos/src/openai-server-compaction/tests/session.effect.test.ts packages/agentos/src/mate-memory/tests/extension.effect.test.ts packages/agentos/src/mate-memory/tests/maintenance.effect.test.ts packages/agentos/src/mate-memory/tests/model.effect.test.ts packages/agentos/src/workloads/tests/compiler.effect.test.ts packages/agentos/src/workloads/tests/compiler-kustomize.effect.test.ts packages/agentos/src/workloads/tests/profiles.effect.test.ts packages/agentos/src/workloads/tests/recovery-conformance.effect.test.ts packages/agentos/src/workloads/tests/disposable-recovery.effect.test.ts packages/agentos/runtime/tests/crewmate-readiness.effect.test.ts packages/agentos/runtime/tests/health.effect.test.ts packages/agentos/runtime/tests/readiness.effect.test.ts packages/agentos/runtime/tests/github-provider.effect.test.ts packages/agentos/runtime/tests/github-workload-auth.effect.test.ts packages/agentos/runtime/tests/memory-activity.effect.test.ts packages/agentos/runtime/tests/memory-store.effect.test.ts packages/agentos/runtime/tests/production-dependencies.effect.test.ts services/a2a/tests services/agentgateway/tests/contract.effect.test.ts services/ai-gateway/tests/accounts.effect.test.ts services/ai-gateway/tests/attribution.effect.test.ts services/ai-gateway/tests/cli.effect.test.ts services/ai-gateway/tests/codex-oauth.effect.test.ts services/ai-gateway/tests/config.effect.test.ts services/ai-gateway/tests/forward.effect.test.ts services/ai-gateway/tests/gateway-service.effect.test.ts services/ai-gateway/tests/kubernetes.effect.test.ts services/ai-gateway/tests/main.effect.test.ts services/ai-gateway/tests/managed-account-live.effect.test.ts services/ai-gateway/tests/observability.effect.test.ts services/ai-gateway/tests/otlp.effect.test.ts services/ai-gateway/tests/selection.effect.test.ts services/ai-gateway/tests/telemetry.effect.test.ts services/ai-gateway/tests/telemetry-overhead.effect.test.ts services/ai-gateway/tests/provider-http.effect.test.ts services/ai-gateway/tests/quota.effect.test.ts services/ai-gateway/tests/response-usage.effect.test.ts services/ai-gateway/tests/state.effect.test.ts services/ai-gateway/tests/storage.effect.test.ts services/egress-authz/tests services/github-broker/tests", + "effect:test": "vitest run --no-file-parallelism mise.effect.test.ts tooling/effect-migration clis/github-app-token/tests/github-app-token.effect.test.ts packages/agentos/tests/*.effect.test.ts services/openfga/tests/*.effect.test.ts benchmarks/tests/contracts.effect.test.ts benchmarks/tests/pi-session-adapter.effect.test.ts benchmarks/tests/runner.effect.test.ts database/kubernetes/cloudnative-pg/tests/manifest.effect.test.ts database/tests/*.effect.test.ts packages/agentos/resources/crewmates/default/kubernetes/tests/manifest.effect.test.ts packages/agentos/resources/roles/firstmate/kubernetes/tests/manifests.effect.test.ts packages/agentos/resources/roles/secondmate/kubernetes/tests/*.effect.test.ts packages/agentos/runtime/kubernetes/tests/*.effect.test.ts packages/agentos/runtime/tests/kubernetes.effect.test.ts services/otel-collector/kubernetes/tests/*.effect.test.ts services/otel-collector/tests/*.effect.test.ts packages/agentos/src/shared/tests packages/agentos/src/access/tests packages/agentos/src/harness-control/tests packages/agentos/src/protocol/tests packages/agentos/src/resilience/tests packages/agentos/src/topology/tests packages/agentos/src/telemetry/tests packages/agentos/src/openai-server-compaction/tests/extension.effect.test.ts packages/agentos/src/openai-server-compaction/tests/messages.effect.test.ts packages/agentos/src/openai-server-compaction/tests/remote-runtime.effect.test.ts packages/agentos/src/openai-server-compaction/tests/schemas.effect.test.ts packages/agentos/src/openai-server-compaction/tests/session.effect.test.ts packages/agentos/src/mate-memory/tests/extension.effect.test.ts packages/agentos/src/mate-memory/tests/maintenance.effect.test.ts packages/agentos/src/mate-memory/tests/model.effect.test.ts packages/agentos/src/workloads/tests/compiler.effect.test.ts packages/agentos/src/workloads/tests/compiler-kustomize.effect.test.ts packages/agentos/src/workloads/tests/profiles.effect.test.ts packages/agentos/src/workloads/tests/recovery-conformance.effect.test.ts packages/agentos/src/workloads/tests/disposable-recovery.effect.test.ts packages/agentos/runtime/tests/crewmate-readiness.effect.test.ts packages/agentos/runtime/tests/health.effect.test.ts packages/agentos/runtime/tests/readiness.effect.test.ts packages/agentos/runtime/tests/github-provider.effect.test.ts packages/agentos/runtime/tests/github-workload-auth.effect.test.ts packages/agentos/runtime/tests/memory-activity.effect.test.ts packages/agentos/runtime/tests/memory-store.effect.test.ts packages/agentos/runtime/tests/production-dependencies.effect.test.ts services/a2a/tests services/agentgateway/tests/contract.effect.test.ts services/ai-gateway/tests/accounts.effect.test.ts services/ai-gateway/tests/attribution.effect.test.ts services/ai-gateway/tests/cli.effect.test.ts services/ai-gateway/tests/codex-oauth.effect.test.ts services/ai-gateway/tests/config.effect.test.ts services/ai-gateway/tests/forward.effect.test.ts services/ai-gateway/tests/gateway-service.effect.test.ts services/ai-gateway/tests/kubernetes.effect.test.ts services/ai-gateway/tests/main.effect.test.ts services/ai-gateway/tests/managed-account-live.effect.test.ts services/ai-gateway/tests/observability.effect.test.ts services/ai-gateway/tests/otlp.effect.test.ts services/ai-gateway/tests/selection.effect.test.ts services/ai-gateway/tests/telemetry.effect.test.ts services/ai-gateway/tests/telemetry-overhead.effect.test.ts services/ai-gateway/tests/provider-http.effect.test.ts services/ai-gateway/tests/workload-client-proxy.effect.test.ts services/ai-gateway/tests/quota.effect.test.ts services/ai-gateway/tests/response-usage.effect.test.ts services/ai-gateway/tests/state.effect.test.ts services/ai-gateway/tests/storage.effect.test.ts services/egress-authz/tests services/github-broker/tests", "effect:runtime:test": "vitest run --no-file-parallelism packages/agentos/src/background-tasks/tests packages/agentos/src/coordination-readiness/tests packages/agentos/src/supervision-guard/tests/extension.effect.test.ts", "effect:bun:test": "bun ./node_modules/vitest/vitest.mjs run --no-file-parallelism tooling/testing/tests/bun-http.effect.test.ts packages/agentos/runtime/tests/capacity-preflight.effect.test.ts packages/agentos/runtime/tests/codex-otel.effect.test.ts packages/agentos/runtime/tests/codex-native-otel.effect.test.ts packages/agentos/runtime/tests/codex-provider.effect.test.ts packages/agentos/runtime/tests/codex-token.effect.test.ts packages/agentos/runtime/tests/distribution-migration.effect.test.ts packages/agentos/runtime/tests/github-native-clients.effect.test.ts packages/agentos/runtime/tests/image-seed.effect.test.ts packages/agentos/runtime/tests/pi-observability-control.effect.test.ts packages/agentos/runtime/tests/pi-provider.effect.test.ts packages/agentos/runtime/tests/prepare-home.effect.test.ts packages/agentos/runtime/tests/readiness-state.effect.test.ts packages/agentos/runtime/tests/runtime.effect.test.ts services/agentgateway/tests/a2a-v1.effect.test.ts services/agentgateway/tests/conformance.effect.test.ts services/ai-gateway/tests/router-package.effect.test.ts services/ai-gateway/tests/routing-state.effect.test.ts && bun run effect:runtime:test", "check": "bun run effect:check && bun run effect:tooling:test && bun run test && bun run typecheck && bun run site:lint && bun run site:typecheck", diff --git a/packages/agentos/skills/agentos-ai-gateway/SKILL.md b/packages/agentos/skills/agentos-ai-gateway/SKILL.md index d4e80abb..632d626c 100644 --- a/packages/agentos/skills/agentos-ai-gateway/SKILL.md +++ b/packages/agentos/skills/agentos-ai-gateway/SKILL.md @@ -221,6 +221,41 @@ a real process/Pod restart; Pi `/reload` cannot change environment. Ask before interrupting a Mate and preserve its native session reference through the normal recovery procedure. +## Hermes workload-proxy migration + +Before any Hermes migration, read the [Hermes integration +contract](../../../../docs/integrations/hermes-ai-gateway.md) and verify that +the selected pinned revision passes its deployment gate. Do not deploy or +update this integration while that gate is unmet; auxiliary requests remain out +of scope until the selected revision satisfies the contract's no-retry, +no-redirect and no-fallback requirements. + +For the live v0.1.24 Hermes topology, the observed pre-migration state uses one +`ai-gateway` replica with a shared `AI_GATEWAY_TOKEN`, while Hermes uses direct +in-Pod provider authentication and has no projected egress identity. Do not +mutate it in place. + +1. Keep Hermes on verified direct OAuth while deploying the reviewed + Agentgateway plus `agentos-egress-authz` topology and registering the Hermes + workload identity/access profile. +2. Build and publish the reviewed AgentOS revision, then update the Hermes + manifest with the exact immutable adapter image digest and projected token + wiring. Do not copy the legacy shared token or any `ai-gateway-client` Secret. +3. Render and review the StatefulSet and NetworkPolicies. Confirm the proxy is + loopback-only, Agentgateway is the only AI Gateway ingress, and ordinary + Internet egress is unchanged. +4. After explicit rollout approval, start a new Hermes session with the exact + selected model and authorize one short fixed no-tool response. Verify the + effective provider/model and failure fidelity without reading prompts, + responses, tokens, account IDs, or vault state. +5. Keep direct OAuth intact until that evidence is accepted. Roll back by + restoring the prior Hermes config/workload revision; do not delete provider + auth or retained homes as part of route rollback. +6. Retire the legacy shared-token request path only after every selected client + has either passed the workload-identity route or returned to verified direct + authentication. Secret removal, deployment, restart, and provider login are + separate human approval gates. + For First and Second Mate, the additive client patch sets `AGENTOS_PI_PROVIDER_MODE=ai-gateway` on `prepare-home`. Before Pi can start, AgentOS atomically reconciles only its marker-owned `openai-codex` transport diff --git a/services/ai-gateway/README.md b/services/ai-gateway/README.md index 965a83ce..330289e4 100644 --- a/services/ai-gateway/README.md +++ b/services/ai-gateway/README.md @@ -59,6 +59,13 @@ composes only the approved clients. Repeated Gateway-owned device logins create the multi-subscription pool; `$agentos-ai-gateway` owns the exact lifecycle and native Pi/Codex configuration. +Responses clients that cannot reread a projected Kubernetes token on every +request use the loopback-only `ai-gateway-workload-proxy` sidecar. Its exact +Responses paths, identity and failure contract, timeout settings, health +checks, Pod wiring, and migration from the legacy shared-token topology are +documented in +[`docs/integrations/hermes-ai-gateway.md`](../../docs/integrations/hermes-ai-gateway.md). + `GET /healthz` proves only that the process is serving. `GET /readyz` requires both usable pooled/provider fallback capacity and the current `agentos-provider-budget-settlement` Pod token to pass the authorizer's diff --git a/services/ai-gateway/package.json b/services/ai-gateway/package.json index 513a7a1b..92000fd2 100644 --- a/services/ai-gateway/package.json +++ b/services/ai-gateway/package.json @@ -5,7 +5,8 @@ "private": true, "type": "module", "bin": { - "ai-gateway": "./src/main.ts" + "ai-gateway": "./src/main.ts", + "ai-gateway-workload-proxy": "./src/workload-client-proxy-main.ts" }, "scripts": { "test": "bun ../../node_modules/vitest/vitest.mjs run tests", diff --git a/services/ai-gateway/src/config.ts b/services/ai-gateway/src/config.ts index 5d5f8d58..2a392bc0 100644 --- a/services/ai-gateway/src/config.ts +++ b/services/ai-gateway/src/config.ts @@ -9,6 +9,9 @@ import { Schema, } from "effect"; +export const defaultAIGatewayIdleTimeoutSeconds = 255; +export const defaultAIGatewayGracefulShutdownMillis = 20_000; + const RawConfig = Config.all({ home: Config.string("HOME").pipe(Config.withDefault(".")), stateDirectory: Config.string("AI_GATEWAY_STATE_DIR").pipe( @@ -21,11 +24,11 @@ const RawConfig = Config.all({ Config.withDefault(8787), ), idleTimeoutSeconds: Config.int("AI_GATEWAY_IDLE_TIMEOUT_SECONDS").pipe( - Config.withDefault(255), + Config.withDefault(defaultAIGatewayIdleTimeoutSeconds), ), gracefulShutdownMillis: Config.int( "AI_GATEWAY_GRACEFUL_SHUTDOWN_MILLIS", - ).pipe(Config.withDefault(20_000)), + ).pipe(Config.withDefault(defaultAIGatewayGracefulShutdownMillis)), clientAuthenticationMode: Config.string( "AI_GATEWAY_CLIENT_AUTH_MODE", ).pipe(Config.withDefault("shared_token")), diff --git a/services/ai-gateway/src/effect-routing-state.ts b/services/ai-gateway/src/effect-routing-state.ts index 6b078d6e..d793a3aa 100644 --- a/services/ai-gateway/src/effect-routing-state.ts +++ b/services/ai-gateway/src/effect-routing-state.ts @@ -11,6 +11,7 @@ import { Layer, Option, Ref, + Schedule, } from "effect"; import { @@ -34,11 +35,12 @@ type LastSelection = NonNullable; export function makeEffectAIRoutingStateLayer( path: string, config: RoutingConfig, -) { - const infrastructure = sqliteRoutingStateLayer( + routingLayer: Layer.Layer = + sqliteRoutingStateLayer( path, toRouterConfig(config), - ); + ), +) { return Layer.effect( AIRoutingState, Effect.gen(function*() { @@ -46,6 +48,11 @@ export function makeEffectAIRoutingStateLayer( const lastSelection = yield* Ref.make>( Option.none(), ); + const releaseRoutingLease = (leaseToken: LeaseToken) => + routing.release(leaseToken).pipe( + Effect.asVoid, + Effect.retry({ schedule: Schedule.spaced("1 second"), times: 3 }), + ); const summary: AIRoutingState["Service"]["summary"] = (now) => routeEffect(Effect.gen(function*() { @@ -67,9 +74,9 @@ export function makeEffectAIRoutingStateLayer( }; })); - const acquire: AIRoutingState["Service"]["acquire"] = (input) => - routeEffect(Effect.gen(function*() { - let durable = yield* routing.summary(input.now); + const acquire: AIRoutingState["Service"]["acquire"] = (input, use) => + Effect.gen(function*() { + let durable = yield* routeEffect(routing.summary(input.now)); const recovered = durable.accounts.filter((account) => account.requiresReauthentication && input.candidates.some((candidate) => @@ -77,7 +84,7 @@ export function makeEffectAIRoutingStateLayer( !candidate.needsReauth ) ); - yield* Effect.forEach( + yield* routeEffect(Effect.forEach( recovered, (account) => routing.recordResponse( @@ -86,46 +93,70 @@ export function makeEffectAIRoutingStateLayer( input.now, ), { discard: true }, - ); + )); if (recovered.length > 0) { - durable = yield* routing.summary(input.now); + durable = yield* routeEffect(routing.summary(input.now)); } const diagnosticCandidates = overlayRoutingSummary( input.candidates, durable, ); - const decision = yield* selectEffect( + const decision = yield* routeEffect(selectEffect( diagnosticCandidates, config, input.now, + )); + const transferred = yield* Ref.make(false); + return yield* Effect.acquireUseRelease( + routing.acquire({ + candidates: input.candidates.map(toRouterCandidate), + now: input.now, + ...(input.sessionKey === undefined + ? {} + : { sessionKey: SessionKey.make(input.sessionKey) }), + }).pipe( + Effect.mapError(() => + AIRoutingStateError.make({ code: "state_unavailable" }) + ), + ), + (lease) => Effect.gen(function*() { + const acquired = Option.getOrUndefined(lease); + const decisionReason = acquired === undefined + ? decision.reason + : acquired.accountId === decision.accountId + ? decision.reason + : "current_account_hysteresis"; + yield* Ref.set(lastSelection, Option.some({ + observedAt: input.now, + reason: decisionReason, + candidates: decision.candidates, + })); + const transfer = Effect.uninterruptible( + Ref.set(transferred, true), + ); + return yield* use( + acquired === undefined + ? undefined + : { + accountId: acquired.accountId, + leaseToken: acquired.leaseToken, + expiresAt: acquired.expiresAt, + decisionReason, + }, + transfer, + ); + }), + (lease, exit) => Effect.uninterruptible(Effect.gen(function*() { + const acquired = Option.getOrUndefined(lease); + if (acquired === undefined) return; + const wasTransferred = yield* Ref.get(transferred); + if (wasTransferred && exit._tag === "Success") return; + yield* routeEffect(releaseRoutingLease(acquired.leaseToken)).pipe( + Effect.catchCause(() => Effect.void), + ); + })), ); - const lease = yield* routing.acquire({ - candidates: input.candidates.map(toRouterCandidate), - now: input.now, - ...(input.sessionKey === undefined - ? {} - : { sessionKey: SessionKey.make(input.sessionKey) }), - }); - const acquired = Option.getOrUndefined(lease); - const decisionReason = acquired === undefined - ? decision.reason - : acquired.accountId === decision.accountId - ? decision.reason - : "current_account_hysteresis"; - yield* Ref.set(lastSelection, Option.some({ - observedAt: input.now, - reason: decisionReason, - candidates: decision.candidates, - })); - return acquired === undefined - ? undefined - : { - accountId: acquired.accountId, - leaseToken: acquired.leaseToken, - expiresAt: acquired.expiresAt, - decisionReason, - }; - })); + }); const evaluate: AIRoutingState["Service"]["evaluate"] = (input) => routeEffect(Effect.gen(function*() { @@ -148,7 +179,7 @@ export function makeEffectAIRoutingStateLayer( renew: (leaseToken, now) => routeEffect(routing.renew(LeaseToken.make(leaseToken), now)), release: (leaseToken) => - routeEffect(routing.release(LeaseToken.make(leaseToken))).pipe( + routeEffect(releaseRoutingLease(LeaseToken.make(leaseToken))).pipe( Effect.as(true), ), recordResponse: (accountId, status, headers, now) => @@ -159,7 +190,7 @@ export function makeEffectAIRoutingStateLayer( )), }); }), - ).pipe(Layer.provide(infrastructure)); + ).pipe(Layer.provide(routingLayer)); } function overlayRoutingSummary( diff --git a/services/ai-gateway/src/forward.ts b/services/ai-gateway/src/forward.ts index 7cbc7e86..24889a29 100644 --- a/services/ai-gateway/src/forward.ts +++ b/services/ai-gateway/src/forward.ts @@ -18,6 +18,7 @@ import { Effect, Exit, Option, + Ref, Result, Schema, Stream, @@ -81,12 +82,16 @@ export type AIForwardClientAuthentication = export interface AIForwardOptions { readonly authentication: AIForwardClientAuthentication; - readonly acquire: ( + readonly acquire: ( sessionKey: string | undefined, signal: AbortSignal, authorization: ProviderAuthorizationGrantV1 | undefined, telemetry: AIGatewayRequestTelemetry, - ) => Effect.Effect; + use: ( + lease: AIForwardLease | undefined, + transfer: Effect.Effect, + ) => Effect.Effect, + ) => Effect.Effect; readonly provider: AIProviderHttp["Service"]; readonly settlements: ProviderBudgetSettlementReporter["Service"]; readonly now: Effect.Effect; @@ -188,180 +193,213 @@ export const makeAIForwardHandler = Effect.fn( } const sessionKey = attributed.success; yield* diagnostic(requestTelemetry.routeStarted); - const acquired = yield* Effect.result(options.acquire( + return yield* options.acquire( sessionKey, request.signal, authentication.authorization, requestTelemetry, - )); - if (Result.isFailure(acquired)) { - yield* diagnostic( - requestTelemetry.routeEnded("error", acquired.failure), - ); - yield* diagnostic(requestTelemetry.end({ - status: 503, - error: acquired.failure, - streamOutcome: "not_streamed", - })); - return jsonResponse(503, "route_unavailable"); - } - const lease = acquired.success; - if (lease === undefined) { - yield* diagnostic(requestTelemetry.routeEnded("unavailable")); - yield* diagnostic(requestTelemetry.end({ - status: 503, - streamOutcome: "not_streamed", - })); - return jsonResponse(503, "no_eligible_account"); - } - yield* diagnostic(requestTelemetry.routeEnded("acquired")); - - const upstreamRequest = yield* makeUpstreamRequest(request, url, lease); - if (upstreamRequest === undefined) { - yield* releaseLease(lease, requestTelemetry); - yield* diagnostic(requestTelemetry.end({ - status: 400, - streamOutcome: "not_streamed", - })); - return jsonResponse(400, "invalid_request"); - } - yield* diagnostic( - requestTelemetry.upstreamStarted(upstreamRequest.headers), - ); - const upstreamResult = yield* Effect.result( - options.provider.execute(upstreamRequest), - ); - if (Result.isFailure(upstreamResult)) { - yield* diagnostic( - requestTelemetry.upstreamFailed(upstreamResult.failure), - ); - yield* releaseLease(lease, requestTelemetry); - yield* diagnostic(requestTelemetry.end({ - status: 502, - error: upstreamResult.failure, - streamOutcome: request.signal.aborted ? "aborted" : "upstream_error", - })); - return jsonResponse(502, providerErrorCode(upstreamResult.failure)); - } - const upstream = upstreamResult.success; - const headersResult = yield* Effect.result(Effect.try({ - try: () => new Headers(upstream.headers), - catch: () => AIForwardConfigurationError.make({ - code: "invalid_configuration", - }), - })); - if (Result.isFailure(headersResult)) { - yield* diagnostic( - requestTelemetry.upstreamFailed(headersResult.failure), - ); - yield* releaseLease(lease, requestTelemetry); - yield* diagnostic(requestTelemetry.end({ - status: 502, - error: headersResult.failure, - streamOutcome: "upstream_error", - })); - return jsonResponse(502, "invalid_provider_response"); - } - const upstreamHeaders = headersResult.success; - yield* diagnostic( - requestTelemetry.upstreamHeaders(upstream.status, upstreamHeaders), - ); - const responseHeaders = sanitizeResponseHeaders(upstreamHeaders); - const responseMetadata = yield* Effect.result(finiteResponse( - null, - upstream.status, - responseHeaders, - )); - if (Result.isFailure(responseMetadata)) { - yield* diagnostic( - requestTelemetry.upstreamFailed(responseMetadata.failure), - ); - yield* releaseLease(lease, requestTelemetry); - yield* diagnostic(requestTelemetry.end({ - status: 502, - error: responseMetadata.failure, - streamOutcome: "upstream_error", - })); - return jsonResponse(502, "invalid_provider_response"); - } - yield* recordResponse(lease, upstream.status, upstreamHeaders); - if (upstream.body === null) { - yield* settleWithoutBody( - options.settlements, - authentication.authorization, - upstream.status, - ); - yield* releaseLease(lease, requestTelemetry); - yield* diagnostic(requestTelemetry.end({ - status: upstream.status, - streamOutcome: "not_streamed", - })); - return responseMetadata.success; - } - if (!responseStatusAllowsBody(upstream.status)) { - const failure = AIForwardConfigurationError.make({ - code: "invalid_configuration", - }); - yield* diagnostic(requestTelemetry.upstreamFailed(failure)); - yield* releaseLease(lease, requestTelemetry); - yield* diagnostic(requestTelemetry.end({ - status: 502, - error: failure, - streamOutcome: "upstream_error", - })); - return jsonResponse(502, "invalid_provider_response"); - } - - const observerResult = upstream.status < 400 && - isEventStream(upstreamHeaders) - ? yield* Effect.result(makeOpenAITerminalUsageObserver({ - maximumEventBytes: options.maximumUsageEventBytes, - })) - : undefined; - if (observerResult !== undefined && Result.isFailure(observerResult)) { - yield* releaseLease(lease, requestTelemetry); - return jsonResponse(503, "accounting_unavailable"); - } - const observer = observerResult === undefined - ? undefined - : observerResult.success; - const monitored = upstream.body.pipe( - Stream.tap((chunk) => - diagnostic(requestTelemetry.streamChunk(chunk.byteLength)) - ), - observer === undefined - ? (stream) => stream - : Stream.tap((chunk) => observer.observe(chunk)), - Stream.onExit((exit) => - finalizeStream( - exit, - lease, - options.settlements, - authentication.authorization, - upstream.status, - observer, - requestTelemetry, - request.signal, - ) + (lease, transfer) => { + if (lease === undefined) { + return Effect.gen(function*() { + yield* diagnostic(requestTelemetry.routeEnded("unavailable")); + yield* diagnostic(requestTelemetry.end({ + status: 503, + streamOutcome: "not_streamed", + })); + return jsonResponse(503, "no_eligible_account"); + }); + } + return Effect.gen(function*() { + type LeaseReleaseState = "available" | "releasing" | "released"; + const leaseReleaseState = yield* Ref.make("available"); + const transferLease = Effect.uninterruptible(transfer); + const releaseLeaseOnce = Effect.uninterruptible(Effect.gen(function*() { + const action = yield* Ref.modify( + leaseReleaseState, + (state): readonly ["attempt" | "done" | "wait", LeaseReleaseState] => + state === "available" + ? ["attempt", "releasing"] + : state === "released" + ? ["done", state] + : ["wait", state], + ); + if (action === "done") return true; + if (action === "wait") return false; + const released = yield* releaseLease(lease, requestTelemetry); + if (released) { + yield* Ref.set(leaseReleaseState, "released"); + yield* transferLease; + return true; + } + yield* Ref.set(leaseReleaseState, "available"); + return false; + })); + + yield* diagnostic(requestTelemetry.routeEnded("acquired")); + const upstreamRequest = yield* makeUpstreamRequest(request, url, lease); + if (upstreamRequest === undefined) { + yield* releaseLeaseOnce; + yield* diagnostic(requestTelemetry.end({ + status: 400, + streamOutcome: "not_streamed", + })); + return jsonResponse(400, "invalid_request"); + } + yield* diagnostic( + requestTelemetry.upstreamStarted(upstreamRequest.headers), + ); + const upstreamResult = yield* Effect.result( + options.provider.execute(upstreamRequest), + ); + if (Result.isFailure(upstreamResult)) { + yield* diagnostic( + requestTelemetry.upstreamFailed(upstreamResult.failure), + ); + yield* releaseLeaseOnce; + yield* diagnostic(requestTelemetry.end({ + status: 502, + error: upstreamResult.failure, + streamOutcome: request.signal.aborted ? "aborted" : "upstream_error", + })); + return jsonResponse(502, providerErrorCode(upstreamResult.failure)); + } + const upstream = upstreamResult.success; + const headersResult = yield* Effect.result(Effect.try({ + try: () => new Headers(upstream.headers), + catch: () => AIForwardConfigurationError.make({ + code: "invalid_configuration", + }), + })); + if (Result.isFailure(headersResult)) { + yield* diagnostic( + requestTelemetry.upstreamFailed(headersResult.failure), + ); + yield* releaseLeaseOnce; + yield* diagnostic(requestTelemetry.end({ + status: 502, + error: headersResult.failure, + streamOutcome: "upstream_error", + })); + return jsonResponse(502, "invalid_provider_response"); + } + const upstreamHeaders = headersResult.success; + yield* diagnostic( + requestTelemetry.upstreamHeaders(upstream.status, upstreamHeaders), + ); + const responseHeaders = sanitizeResponseHeaders(upstreamHeaders); + const responseMetadata = yield* Effect.result(finiteResponse( + null, + upstream.status, + responseHeaders, + )); + if (Result.isFailure(responseMetadata)) { + yield* diagnostic( + requestTelemetry.upstreamFailed(responseMetadata.failure), + ); + yield* releaseLeaseOnce; + yield* diagnostic(requestTelemetry.end({ + status: 502, + error: responseMetadata.failure, + streamOutcome: "upstream_error", + })); + return jsonResponse(502, "invalid_provider_response"); + } + yield* recordResponse(lease, upstream.status, upstreamHeaders); + if (upstream.body === null) { + yield* settleWithoutBody( + options.settlements, + authentication.authorization, + upstream.status, + ); + yield* releaseLeaseOnce; + yield* diagnostic(requestTelemetry.end({ + status: upstream.status, + streamOutcome: "not_streamed", + })); + return responseMetadata.success; + } + if (!responseStatusAllowsBody(upstream.status)) { + const failure = AIForwardConfigurationError.make({ + code: "invalid_configuration", + }); + yield* diagnostic(requestTelemetry.upstreamFailed(failure)); + yield* releaseLeaseOnce; + yield* diagnostic(requestTelemetry.end({ + status: 502, + error: failure, + streamOutcome: "upstream_error", + })); + return jsonResponse(502, "invalid_provider_response"); + } + + const observerResult = upstream.status < 400 && + isEventStream(upstreamHeaders) + ? yield* Effect.result(makeOpenAITerminalUsageObserver({ + maximumEventBytes: options.maximumUsageEventBytes, + })) + : undefined; + if (observerResult !== undefined && Result.isFailure(observerResult)) { + yield* releaseLeaseOnce; + return jsonResponse(503, "accounting_unavailable"); + } + const observer = observerResult === undefined + ? undefined + : observerResult.success; + const monitored = upstream.body.pipe( + Stream.tap((chunk) => + diagnostic(requestTelemetry.streamChunk(chunk.byteLength)) + ), + observer === undefined + ? (stream) => stream + : Stream.tap((chunk) => observer.observe(chunk)), + Stream.onExit((exit) => + finalizeStream( + exit, + releaseLeaseOnce, + options.settlements, + authentication.authorization, + upstream.status, + observer, + requestTelemetry, + request.signal, + ) + ), + ); + const scoped = Stream.unwrap(Effect.gen(function*() { + yield* heartbeat(lease, options.heartbeatMillis).pipe( + Effect.forkScoped({ startImmediately: true }), + ); + return monitored; + })); + const body = yield* Stream.toReadableStreamEffect(scoped); + const responseResult = yield* Effect.result(finiteResponse( + body, + upstream.status, + responseHeaders, + )); + if (Result.isFailure(responseResult)) { + yield* releaseLeaseOnce; + return jsonResponse(502, "invalid_provider_response"); + } + return yield* Effect.uninterruptible(Effect.gen(function*() { + yield* transferLease; + return responseResult.success; + })); + }); + }, + ).pipe( + Effect.catchTag("AIForwardRouteError", (failure) => + Effect.gen(function*() { + yield* diagnostic(requestTelemetry.routeEnded("error", failure)); + yield* diagnostic(requestTelemetry.end({ + status: 503, + error: failure, + streamOutcome: "not_streamed", + })); + return jsonResponse(503, "route_unavailable"); + }) ), ); - const scoped = Stream.unwrap(Effect.gen(function*() { - yield* heartbeat(lease, options.heartbeatMillis).pipe( - Effect.forkScoped({ startImmediately: true }), - ); - return monitored; - })); - const body = yield* Stream.toReadableStreamEffect(scoped); - const responseResult = yield* Effect.result(finiteResponse( - body, - upstream.status, - responseHeaders, - )); - if (Result.isFailure(responseResult)) { - yield* releaseLease(lease, requestTelemetry); - return jsonResponse(502, "invalid_provider_response"); - } - return responseResult.success; }); return handler; }); @@ -493,7 +531,7 @@ function heartbeat(lease: AIForwardLease, heartbeatMillis: number) { function finalizeStream( exit: Exit.Exit, - lease: AIForwardLease, + releaseLeaseOnce: Effect.Effect, settlements: ProviderBudgetSettlementReporter["Service"], authorization: ProviderAuthorizationGrantV1 | undefined, status: number, @@ -501,6 +539,23 @@ function finalizeStream( telemetry: AIGatewayRequestTelemetry, signal: AbortSignal, ): Effect.Effect { + const maximumReleaseAttempts = 3; + let attempts = 0; + let released = false; + const releaseStreamLease = Effect.whileLoop({ + while: () => !released && attempts < maximumReleaseAttempts, + body: () => releaseLeaseOnce.pipe( + Effect.flatMap((attempted) => + attempted || attempts + 1 >= maximumReleaseAttempts + ? Effect.succeed(attempted) + : Effect.sleep(1_000).pipe(Effect.as(false)) + ), + ), + step: (attempted) => { + attempts += 1; + released = attempted; + }, + }); return Effect.gen(function*() { if (authorization !== undefined) { if (status >= 400) { @@ -520,7 +575,7 @@ function finalizeStream( } } } - yield* releaseLease(lease, telemetry); + yield* releaseStreamLease; const outcome = telemetryStreamOutcome(exit, signal); const failure = Exit.isFailure(exit) ? Option.getOrUndefined(Cause.findErrorOption(exit.cause)) @@ -531,7 +586,7 @@ function finalizeStream( ...(failure === undefined ? {} : { error: failure }), })); }).pipe( - Effect.catchCause(() => releaseLease(lease, telemetry)), + Effect.catchCause(() => releaseStreamLease), Effect.catchCause(() => Effect.void), Effect.uninterruptible, ); @@ -588,6 +643,7 @@ function releaseLease( ? telemetry.routeReleased : telemetry.routeReleaseFailed, ); + return Exit.isSuccess(release); }).pipe( Effect.uninterruptible, ); diff --git a/services/ai-gateway/src/gateway-service.ts b/services/ai-gateway/src/gateway-service.ts index e9fb860a..b1e4eec3 100644 --- a/services/ai-gateway/src/gateway-service.ts +++ b/services/ai-gateway/src/gateway-service.ts @@ -7,6 +7,7 @@ import { import { Clock, Effect, + Exit, Ref, Result, Schema, @@ -119,12 +120,16 @@ export const makeAIGatewayApplication = Effect.fn( ); const acquire = Effect.fn("agentos.aiGateway.acquireRoute")( - function*( + function*( sessionKey: string | undefined, _signal: AbortSignal, _authorization: ProviderAuthorizationGrantV1 | undefined, telemetry: AIGatewayRequestTelemetry, - ): Effect.fn.Return { + use: ( + lease: AIForwardLease | undefined, + transfer: Effect.Effect, + ) => Effect.Effect, + ): Effect.fn.Return { const set = yield* candidateSet(telemetry); const currentTime = yield* Clock.currentTimeMillis; yield* Effect.forEach( @@ -141,67 +146,99 @@ export const makeAIGatewayApplication = Effect.fn( ).pipe(Effect.catchCause(() => Effect.void)), { discard: true }, ); - const reservation = yield* routing.acquire({ + return yield* routing.acquire({ candidates: set.candidates, now: currentTime, ...(sessionKey === undefined ? {} : { sessionKey }), - }).pipe(Effect.mapError(() => routeError("routing_unavailable"))); - if (reservation === undefined) { - return fallbackAvailable && fallbackKey !== undefined - ? fallbackLease(fallbackKey) - : undefined; - } - const credentialResult = yield* Effect.result( - vault.getFreshCredential(reservation.accountId), - ); - if (Result.isFailure(credentialResult)) { - yield* routing.release(reservation.leaseToken).pipe( - Effect.catchCause(() => Effect.void), - ); - return yield* routeError(routeCodeForAccount(credentialResult.failure)); - } - const credential = credentialResult.success; - const lease: AIForwardLease = { - kind: "codex_oauth", - accessToken: credential.accessToken, - providerAccountId: credential.providerAccountId, - renew: Clock.currentTimeMillis.pipe( - Effect.flatMap((renewedAt) => - routing.renew(reservation.leaseToken, renewedAt) - ), - Effect.mapError(() => routeError("state_unavailable")), - ), - release: routing.release(reservation.leaseToken).pipe( - Effect.asVoid, - Effect.mapError(() => routeError("state_unavailable")), - ), - recordResponse: (status, headers) => - Effect.gen(function*() { - const responseAt = yield* Clock.currentTimeMillis; - if (status === 401) { - yield* vault.markNeedsReauth( - reservation.accountId, - credential.accessToken, - ).pipe(Effect.catchCause(() => Effect.succeed(false))); + }, (reservation, transfer) => { + if (reservation === undefined) { + return use( + fallbackAvailable && fallbackKey !== undefined + ? fallbackLease(fallbackKey) + : undefined, + Effect.void, + ); + } + return Effect.gen(function*() { + const credential = yield* vault.getFreshCredential( + reservation.accountId, + ).pipe( + Effect.mapError((error) => + routeError(routeCodeForAccount(error)) + ), + ); + type ReleaseState = "available" | "releasing" | "released"; + const releaseState = yield* Ref.make("available"); + const release = Effect.uninterruptible(Effect.gen(function*() { + const action = yield* Ref.modify( + releaseState, + (state): readonly ["attempt" | "done" | "wait", ReleaseState] => + state === "available" + ? ["attempt", "releasing"] + : state === "released" + ? ["done", state] + : ["wait", state], + ); + if (action !== "attempt") return; + const released = yield* Effect.exit( + routing.release(reservation.leaseToken).pipe( + Effect.asVoid, + Effect.mapError(() => routeError("state_unavailable")), + ), + ); + if (Exit.isSuccess(released)) { + yield* Ref.set(releaseState, "released"); + return; } - yield* routing.recordResponse( - reservation.accountId, - status, - headers, - responseAt, - ).pipe( + yield* Ref.set(releaseState, "available"); + return yield* Effect.failCause(released.cause); + })); + const lease: AIForwardLease = { + kind: "codex_oauth", + accessToken: credential.accessToken, + providerAccountId: credential.providerAccountId, + renew: Clock.currentTimeMillis.pipe( + Effect.flatMap((renewedAt) => + routing.renew(reservation.leaseToken, renewedAt) + ), Effect.mapError(() => routeError("state_unavailable")), - ); - }), - }; - return lease; + ), + release, + recordResponse: (status, headers) => + Effect.gen(function*() { + const responseAt = yield* Clock.currentTimeMillis; + if (status === 401) { + yield* vault.markNeedsReauth( + reservation.accountId, + credential.accessToken, + ).pipe(Effect.catchCause(() => Effect.succeed(false))); + } + yield* routing.recordResponse( + reservation.accountId, + status, + headers, + responseAt, + ).pipe( + Effect.mapError(() => routeError("state_unavailable")), + ); + }), + }; + return yield* use(lease, transfer); + }); + }).pipe( + Effect.mapError((error) => + error._tag === "AIForwardRouteError" + ? error + : routeError("routing_unavailable") + ), + ); }, ); const forward = yield* makeAIForwardHandler({ authentication: options.authentication, - acquire: (sessionKey, signal, authorization, telemetry) => - acquire(sessionKey, signal, authorization, telemetry), + acquire: (sessionKey, signal, authorization, telemetry, use) => + acquire(sessionKey, signal, authorization, telemetry, use), provider, settlements, now: Clock.currentTimeMillis, diff --git a/services/ai-gateway/src/main.ts b/services/ai-gateway/src/main.ts index 21fe7656..5e7f7662 100755 --- a/services/ai-gateway/src/main.ts +++ b/services/ai-gateway/src/main.ts @@ -60,7 +60,10 @@ import { AIGatewayTelemetry, makeAIGatewayTelemetry, } from "./observability.ts"; -import { AIProviderHttp, AIProviderHttpLive } from "./provider-http.ts"; +import { + AIProviderHttp, + makeAIProviderHttpLive, +} from "./provider-http.ts"; import { AIGatewayOtlpLive } from "./otlp.ts"; import { CodexQuota, makeCodexQuotaLayer } from "./quota.ts"; import { makeEffectManagedAccountVaultLayer } from "./managed-account-live.ts"; @@ -129,6 +132,8 @@ const AIGatewayStatusClientLive = Layer.effect( }), ); +const aiProviderHttpLive = makeAIProviderHttpLive(BunHttpClient.layer); + function acquireAIGatewayTelemetry() { return Effect.gen(function*() { const disabled = yield* Config.boolean("OTEL_SDK_DISABLED").pipe( @@ -163,7 +168,7 @@ function makeAIGatewayRuntimeLive( const crypto = yield* Crypto.Crypto; const fileSystem = yield* FileSystem.FileSystem; const provider = yield* AIProviderHttp.pipe( - Effect.provide(AIProviderHttpLive), + Effect.provide(aiProviderHttpLive), ); const quota = yield* CodexQuota.pipe( Effect.provide(makeCodexQuotaLayer(config.quotaTimeoutMillis)), diff --git a/services/ai-gateway/src/provider-http.ts b/services/ai-gateway/src/provider-http.ts index 97165ded..176561f9 100644 --- a/services/ai-gateway/src/provider-http.ts +++ b/services/ai-gateway/src/provider-http.ts @@ -1,6 +1,7 @@ import { classifyAIError } from "@akua-dev/agentos"; import { Context, Effect, Layer, Schema, Stream } from "effect"; import { + FetchHttpClient, HttpClient, HttpClientError, HttpClientRequest, @@ -46,14 +47,19 @@ export const AIProviderHttpLive = Layer.effect( try: () => HttpClientRequest.fromWeb(request), catch: () => providerHttpError("request_invalid"), }); - const response = yield* client.execute(clientRequest).pipe( + const response = yield* Effect.raceFirst( + client.execute(clientRequest), + abortOnSignal(request.signal), + ).pipe( Effect.mapError((error) => providerHttpError(httpErrorCode(error))), ); const hasNoBody = request.method === "HEAD" || - response.status === 204 || response.status === 304; + response.status === 204 || response.status === 205 || + response.status === 304; const body: AIProviderResponse["body"] = hasNoBody ? null : response.stream.pipe( + Stream.interruptWhen(abortOnSignal(request.signal)), Stream.mapError((error) => providerHttpError(streamErrorCode(error)) ), @@ -69,6 +75,31 @@ export const AIProviderHttpLive = Layer.effect( }), ); +export const AIProviderHttpRequestInit = Layer.succeed( + FetchHttpClient.RequestInit, + { redirect: "manual" }, +); + +export function makeAIProviderHttpLive( + clientLayer: Layer.Layer, +) { + return AIProviderHttpLive.pipe( + Layer.provide(clientLayer.pipe(Layer.provide(AIProviderHttpRequestInit))), + ); +} + +function abortOnSignal(signal: AbortSignal) { + return Effect.callback((resume) => { + const onAbort = () => resume(Effect.interrupt); + if (signal.aborted) { + onAbort(); + return Effect.void; + } + signal.addEventListener("abort", onAbort, { once: true }); + return Effect.sync(() => signal.removeEventListener("abort", onAbort)); + }); +} + function providerHttpError(code: AIProviderHttpError["code"]) { return AIProviderHttpError.make({ code }); } diff --git a/services/ai-gateway/src/routing-service.ts b/services/ai-gateway/src/routing-service.ts index 0ed43fe4..8da726d2 100644 --- a/services/ai-gateway/src/routing-service.ts +++ b/services/ai-gateway/src/routing-service.ts @@ -49,9 +49,13 @@ export class AIRoutingState extends Context.Service< readonly summary: ( now: number, ) => Effect.Effect; - readonly acquire: ( + readonly acquire: ( input: AIRoutingAcquireInput, - ) => Effect.Effect; + use: ( + reservation: AIAcquiredReservation | undefined, + transfer: Effect.Effect, + ) => Effect.Effect, + ) => Effect.Effect; readonly evaluate: ( input: AIRoutingEvaluateInput, ) => Effect.Effect; diff --git a/services/ai-gateway/src/workload-client-proxy-main.ts b/services/ai-gateway/src/workload-client-proxy-main.ts new file mode 100644 index 00000000..c5ecaaf0 --- /dev/null +++ b/services/ai-gateway/src/workload-client-proxy-main.ts @@ -0,0 +1,138 @@ +#!/usr/bin/env bun + +import * as BunFileSystem from "@effect/platform-bun/BunFileSystem"; +import * as BunHttpClient from "@effect/platform-bun/BunHttpClient"; +import * as BunHttpServer from "@effect/platform-bun/BunHttpServer"; +import * as BunRuntime from "@effect/platform-bun/BunRuntime"; +import { + ConfigProvider, + Console, + Effect, + Layer, + Stream, +} from "effect"; +import { + HttpRouter, + HttpServerRequest, + HttpServerResponse, +} from "effect/unstable/http"; + +import { + AIProviderHttp, + makeAIProviderHttpLive, + type AIProviderResponse, +} from "./provider-http.ts"; +import { + loadWorkloadClientProxyConfig, + makeWorkloadClientProxyHandler, + workloadClientProxyReadinessResponse, + workloadClientProxyErrorResponse, + WorkloadClientProxyError, +} from "./workload-client-proxy.ts"; + +const hopByHopResponseHeaders = new Set([ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", +]); +const httpToken = /^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$/; + +function responseHeaders(upstream: AIProviderResponse) { + const headers = new Headers(upstream.headers); + const connection = headers.get("connection"); + headers.delete("connection"); + for (const name of connection?.split(",") ?? []) { + const normalized = name.trim().toLowerCase(); + if (httpToken.test(normalized)) headers.delete(normalized); + } + for (const name of hopByHopResponseHeaders) headers.delete(name); + headers.delete("content-encoding"); + headers.delete("content-length"); + return headers; +} + +export function responseFromUpstream( + upstream: AIProviderResponse, +) { + return Effect.gen(function*() { + const headers = responseHeaders(upstream); + if (upstream.body === null || upstream.status === 205) { + return new Response(null, { status: upstream.status, headers }); + } + const body = yield* Stream.toReadableStreamEffect(upstream.body); + return new Response(body, { status: upstream.status, headers }); + }); +} + +const startup = Effect.gen(function*() { + const config = yield* loadWorkloadClientProxyConfig(); + const http = yield* AIProviderHttp; + const handler = yield* makeWorkloadClientProxyHandler({ + upstreamBaseUrl: config.upstreamBaseUrl, + tokenPath: config.tokenPath, + assignmentId: config.assignmentId, + forward: (request) => + http.execute(request).pipe( + Effect.mapError(() => + WorkloadClientProxyError.make({ code: "upstream_unavailable" }) + ), + Effect.flatMap(responseFromUpstream), + ), + }); + const routes = Layer.effectDiscard(Effect.gen(function*() { + const router = yield* HttpRouter.HttpRouter; + yield* router.add("GET", "/livez", () => + Effect.succeed(HttpServerResponse.fromWeb(Response.json({ status: "alive" })))); + yield* router.add("GET", "/readyz", () => + workloadClientProxyReadinessResponse(config.tokenPath).pipe( + Effect.map(HttpServerResponse.fromWeb), + )); + yield* router.add("*", "/*", (request) => + HttpServerRequest.toWeb(request).pipe( + Effect.flatMap(handler), + Effect.catchTag("WorkloadClientProxyError", (error) => + Effect.succeed(workloadClientProxyErrorResponse(error))), + Effect.catch(() => + Effect.succeed(Response.json({ error: "proxy_unavailable" }, { + status: 503, + })) + ), + Effect.map(HttpServerResponse.fromWeb), + )); + })); + const server = HttpRouter.serve(routes, { disableListenLog: true }).pipe( + Layer.provide(BunHttpServer.layer({ + hostname: config.hostname, + port: config.port, + idleTimeout: config.idleTimeoutSeconds, + gracefulShutdownTimeout: config.gracefulShutdownMillis, + })), + ); + yield* Console.log(JSON.stringify({ + event: "agentos.ai_gateway_workload_proxy.listening", + hostname: config.hostname, + port: config.port, + })); + return yield* Layer.launch(server); +}); + +if (import.meta.main) { + const platform = Layer.mergeAll( + BunFileSystem.layer, + makeAIProviderHttpLive(BunHttpClient.layer), + ConfigProvider.layer(ConfigProvider.fromEnv()), + ); + BunRuntime.runMain(startup.pipe( + Effect.tapError(() => + Console.error(JSON.stringify({ + event: "agentos.ai_gateway_workload_proxy.failed", + })) + ), + Effect.provide(platform), + ), { disableErrorReporting: true }); +} diff --git a/services/ai-gateway/src/workload-client-proxy.ts b/services/ai-gateway/src/workload-client-proxy.ts new file mode 100644 index 00000000..b7a2dd6f --- /dev/null +++ b/services/ai-gateway/src/workload-client-proxy.ts @@ -0,0 +1,226 @@ +import { sanitizeRequestHeaders } from "@akua-dev/codex-router/codex"; +import { Config, Effect, FileSystem, Option, Schema } from "effect"; + +import { + defaultAIGatewayGracefulShutdownMillis, + defaultAIGatewayIdleTimeoutSeconds, +} from "./config.ts"; + +const maximumTokenBytes = 16 * 1024; +const jwtLike = /^[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+$/; +const allowedPaths = new Set(["/v1/responses", "/v1/responses/compact"]); +const reviewedAgentgatewayOrigin = + "http://agentgateway-openai.agentos.svc.cluster.local:8788"; +const assignmentIdPattern = + /^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i; +export const workloadClientProxyHostname = "127.0.0.1"; +export class WorkloadClientProxyError extends Schema.TaggedErrorClass()( + "WorkloadClientProxyError", + { + code: Schema.Literals([ + "invalid_request", + "token_unavailable", + "upstream_unavailable", + ]), + }, +) {} + +export interface WorkloadClientProxyOptions { + readonly upstreamBaseUrl: URL; + readonly tokenPath: string; + readonly assignmentId?: string; + readonly forward: ( + request: Request, + ) => Effect.Effect; +} + +const Port = Schema.Number.pipe( + Schema.check(Schema.isInt(), Schema.isGreaterThan(0)), + Schema.check(Schema.isLessThanOrEqualTo(65_535)), +); +const Configuration = Schema.Struct({ + port: Port, + idleTimeoutSeconds: Schema.Number.pipe( + Schema.check( + Schema.isInt(), + Schema.isGreaterThanOrEqualTo(0), + Schema.isLessThanOrEqualTo(255), + ), + ), + tokenPath: Schema.String.pipe(Schema.check(Schema.isMinLength(1))), + upstreamBaseUrl: Schema.URL, + assignmentId: Schema.Union([ + Schema.Literal(""), + Schema.String.pipe(Schema.check(Schema.isPattern(assignmentIdPattern))), + ]), + gracefulShutdownMillis: Schema.Number.pipe( + Schema.check(Schema.isInt(), Schema.isGreaterThan(0)), + ), +}); + +export class WorkloadClientProxyConfigurationError extends Schema.TaggedErrorClass()( + "WorkloadClientProxyConfigurationError", + { code: Schema.Literal("invalid_configuration") }, +) {} + +function configurationError() { + return WorkloadClientProxyConfigurationError.make({ + code: "invalid_configuration", + }); +} + +function isValidUpstreamBaseUrl(url: URL) { + return url.origin === reviewedAgentgatewayOrigin && + url.username === "" && + url.password === "" && + url.pathname === "/" && + url.search === "" && + url.hash === ""; +} + +export const loadWorkloadClientProxyConfig = Effect.fn( + "agentos.aiGateway.workloadClient.loadConfig", +)(function*() { + const raw = yield* Config.all({ + port: Config.int("AI_GATEWAY_WORKLOAD_PROXY_PORT").pipe( + Config.withDefault(8_790), + ), + idleTimeoutSeconds: Config.int("AI_GATEWAY_IDLE_TIMEOUT_SECONDS").pipe( + Config.withDefault(defaultAIGatewayIdleTimeoutSeconds), + ), + tokenPath: Config.string("AGENTOS_EGRESS_TOKEN_FILE").pipe( + Config.withDefault("/var/run/secrets/agentos-egress/token"), + ), + upstreamBaseUrl: Config.url("AI_GATEWAY_URL"), + assignmentId: Config.string("AGENTOS_ASSIGNMENT_ID").pipe( + Config.withDefault(""), + ), + gracefulShutdownMillis: Config.int("AI_GATEWAY_GRACEFUL_SHUTDOWN_MILLIS").pipe( + Config.withDefault(defaultAIGatewayGracefulShutdownMillis), + ), + }).pipe(Effect.mapError(configurationError)); + const config = yield* Schema.decodeUnknownEffect(Configuration)(raw).pipe( + Effect.mapError(configurationError), + ); + if (!isValidUpstreamBaseUrl(config.upstreamBaseUrl)) { + return yield* configurationError(); + } + return { + ...config, + hostname: workloadClientProxyHostname, + assignmentId: config.assignmentId === "" ? undefined : config.assignmentId, + }; +}); + +function proxyError(code: WorkloadClientProxyError["code"]) { + return WorkloadClientProxyError.make({ code }); +} + +export const readProjectedWorkloadToken = Effect.fn( + "agentos.aiGateway.workloadClient.readToken", +)(function*(path: string) { + const fileSystem = yield* FileSystem.FileSystem; + const bytes = yield* Effect.scoped(Effect.gen(function*() { + const file = yield* fileSystem.open(path, { flag: "r" }); + const read = yield* file.readAlloc(maximumTokenBytes + 1); + return Option.getOrElse(read, () => new Uint8Array()); + })).pipe( + Effect.mapError(() => proxyError("token_unavailable")), + ); + if (bytes.length === 0 || bytes.length > maximumTokenBytes) { + return yield* proxyError("token_unavailable"); + } + const token = yield* Effect.try({ + try: () => new TextDecoder("utf-8", { fatal: true }).decode(bytes), + catch: () => proxyError("token_unavailable"), + }); + if (token.trim() !== token || !jwtLike.test(token)) { + return yield* proxyError("token_unavailable"); + } + return token; +}); + +function forwardedHeaders( + input: Headers, + token: string, + assignmentId: string | undefined, +) { + const headers = sanitizeRequestHeaders(input); + for (const name of Array.from(headers.keys())) { + if (name.toLowerCase().startsWith("x-agentos-")) { + headers.delete(name); + } + } + headers.set("authorization", `Bearer ${token}`); + if (assignmentId !== undefined) { + headers.set("x-agentos-assignment-id", assignmentId); + } + return headers; +} + +export const workloadClientProxyReadinessResponse = Effect.fn( + "agentos.aiGateway.workloadClient.readiness", +)(function*(tokenPath: string) { + return yield* readProjectedWorkloadToken(tokenPath).pipe( + Effect.as(Response.json({ status: "ready" })), + Effect.catchTag("WorkloadClientProxyError", () => + Effect.succeed(Response.json({ status: "not_ready" }, { status: 503 }))), + ); +}); + +export const makeWorkloadClientProxyHandler = Effect.fn( + "agentos.aiGateway.workloadClient.makeHandler", +)(function*(options: WorkloadClientProxyOptions) { + if (!isValidUpstreamBaseUrl(options.upstreamBaseUrl)) { + return yield* proxyError("invalid_request"); + } + const handle = Effect.fn("agentos.aiGateway.workloadClient.forward")( + function*(request: Request) { + const incoming = yield* Effect.try({ + try: () => new URL(request.url), + catch: () => proxyError("invalid_request"), + }); + if ( + request.method !== "POST" || + !allowedPaths.has(incoming.pathname) || + incoming.search !== "" || + incoming.hash !== "" + ) { + return yield* proxyError("invalid_request"); + } + const token = yield* readProjectedWorkloadToken(options.tokenPath); + const upstream = new URL(incoming.pathname, options.upstreamBaseUrl); + const forwarded = yield* Effect.try({ + try: () => { + const init: RequestInit & { readonly duplex: "half" } = { + method: "POST", + headers: forwardedHeaders( + request.headers, + token, + options.assignmentId, + ), + body: request.body, + redirect: "manual", + signal: request.signal, + duplex: "half", + }; + return new Request(upstream.toString(), init); + }, + catch: () => proxyError("invalid_request"), + }); + return yield* options.forward(forwarded); + }, + ); + return handle; +}); + +export function workloadClientProxyErrorResponse( + error: WorkloadClientProxyError, +) { + const status = error.code === "invalid_request" + ? 404 + : error.code === "token_unavailable" + ? 503 + : 502; + return Response.json({ error: error.code }, { status }); +} diff --git a/services/ai-gateway/tests/forward.effect.test.ts b/services/ai-gateway/tests/forward.effect.test.ts index c752b402..614c818f 100644 --- a/services/ai-gateway/tests/forward.effect.test.ts +++ b/services/ai-gateway/tests/forward.effect.test.ts @@ -8,7 +8,17 @@ import { type ProviderBudgetSettlementReceiptV1, type ProviderBudgetSettlementReportV1, } from "@akua-dev/agentos"; -import { Effect, Fiber, Layer, Metric, Ref, Stream, Tracer } from "effect"; +import { + Effect, + Exit, + Fiber, + Layer, + Metric, + Option, + Ref, + Stream, + Tracer, +} from "effect"; import { TestClock } from "effect/testing"; import { @@ -96,6 +106,7 @@ const encoder = new TextEncoder(); const makeLease = Effect.fn("test.aiForward.makeLease")(function*() { const releases = yield* Ref.make(0); + const released = yield* Ref.make(false); const renewals = yield* Ref.make(0); const lease: AIForwardLease = { kind: "openai_api_key", @@ -103,12 +114,53 @@ const makeLease = Effect.fn("test.aiForward.makeLease")(function*() { renew: Ref.updateAndGet(renewals, (count) => count + 1).pipe( Effect.as(true), ), - release: Ref.update(releases, (count) => count + 1), + release: Effect.gen(function*() { + const shouldRelease = yield* Ref.modify( + released, + (state): readonly [boolean, boolean] => [!state, true], + ); + if (shouldRelease) yield* Ref.update(releases, (count) => count + 1); + }), recordResponse: () => Effect.void, }; return { lease, releases, renewals }; }); +function acquireLease(lease: AIForwardLease | undefined) { + return ( + _session: string | undefined, + _signal: AbortSignal, + _authorization: ProviderAuthorizationGrantV1 | undefined, + _telemetry: AIGatewayRequestTelemetry, + use: ( + lease: AIForwardLease | undefined, + transfer: Effect.Effect, + ) => Effect.Effect, + ) => { + if (lease === undefined) return use(undefined, Effect.void); + return Effect.gen(function*() { + const transferred = yield* Ref.make(false); + return yield* Effect.acquireUseRelease( + Effect.succeed(lease), + (acquired) => use( + acquired, + Effect.uninterruptible(Ref.set(transferred, true)), + ), + (acquired, exit) => + exit._tag === "Success" && acquired !== undefined + ? Ref.get(transferred).pipe( + Effect.flatMap((wasTransferred) => + wasTransferred ? Effect.void : acquired.release + ), + ) + : acquired === undefined + ? Effect.void + : acquired.release, + ); + }); + }; +} + const makeSettlementRecorder = Effect.fn( "test.aiForward.makeSettlementRecorder", )(function*() { @@ -195,7 +247,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 200, @@ -253,7 +305,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeAIGatewayTelemetry(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 200, @@ -300,7 +352,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 500, headers: {}, body: null }), }), @@ -341,7 +393,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "shared_token", token: "client-secret" }, - acquire: () => Effect.succeed(undefined), + acquire: acquireLease(undefined), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 500, headers: {}, body: null }), }), @@ -385,7 +437,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(undefined), + acquire: acquireLease(undefined), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 500, headers: {}, body: null }), }), @@ -467,7 +519,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(undefined), + acquire: acquireLease(undefined), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 500, headers: {}, body: null }), }), @@ -514,7 +566,7 @@ describe("Effect AI Gateway forwarding", () => { }; const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(lease), + acquire: acquireLease(lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 500, headers: {}, body: null }), }), @@ -572,8 +624,10 @@ describe("Effect AI Gateway forwarding", () => { }); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: (_session, _signal, authorization) => - Ref.set(attribution, authorization).pipe(Effect.as(route.lease)), + acquire: (_session, _signal, authorization, _telemetry, use) => + Ref.set(attribution, authorization).pipe( + Effect.flatMap(() => use(route.lease, Effect.void)), + ), provider, settlements: settlement.settlements, now: Effect.succeed(now), @@ -624,7 +678,7 @@ describe("Effect AI Gateway forwarding", () => { const settlement = yield* makeSettlementRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 200, @@ -651,7 +705,7 @@ describe("Effect AI Gateway forwarding", () => { const settlement = yield* makeSettlementRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 429, @@ -712,7 +766,7 @@ describe("Effect AI Gateway forwarding", () => { }; const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status, @@ -779,7 +833,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 204, @@ -821,6 +875,56 @@ describe("Effect AI Gateway forwarding", () => { ]); })); + it.effect("releases an acquired route when the provider is interrupted before headers", () => + Effect.gen(function*() { + const route = yield* makeLease(); + const settlement = yield* makeSettlementRecorder(); + const handler = yield* makeAIForwardHandler({ + authentication: { kind: "workload_identity" }, + acquire: acquireLease(route.lease), + provider: AIProviderHttp.of({ + execute: () => Effect.interrupt, + }), + settlements: settlement.settlements, + now: Effect.succeed(now), + heartbeatMillis: 40_000, + maximumUsageEventBytes: 4_096, + }); + const exit = yield* Effect.exit(handler(gatewayRequest())); + + assert.isFalse(Exit.isSuccess(exit)); + assert.strictEqual(yield* Ref.get(route.releases), 1); + })); + + it.effect("releases an acquired route when post-acquisition work defects", () => + Effect.gen(function*() { + const route = yield* makeLease(); + const settlement = yield* makeSettlementRecorder(); + const brokenLease: AIForwardLease = { + kind: route.lease.kind, + get accessToken(): string { + return Option.getOrThrow(Option.none()); + }, + renew: route.lease.renew, + release: route.lease.release, + }; + const handler = yield* makeAIForwardHandler({ + authentication: { kind: "workload_identity" }, + acquire: acquireLease(brokenLease), + provider: AIProviderHttp.of({ + execute: () => Effect.die("provider must not execute"), + }), + settlements: settlement.settlements, + now: Effect.succeed(now), + heartbeatMillis: 40_000, + maximumUsageEventBytes: 4_096, + }); + const exit = yield* Effect.exit(handler(gatewayRequest())); + + assert.isFalse(Exit.isSuccess(exit)); + assert.strictEqual(yield* Ref.get(route.releases), 1); + })); + it.effect("ends telemetry when a finite provider response cannot be constructed", () => Effect.gen(function*() { const route = yield* makeLease(); @@ -828,7 +932,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 700, headers: {}, body: null }), }), @@ -875,7 +979,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 700, @@ -926,7 +1030,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 204, @@ -977,9 +1081,10 @@ describe("Effect AI Gateway forwarding", () => { const settlement = yield* makeSettlementRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Ref.update(acquireCalls, (count) => count + 1).pipe( - Effect.as(undefined), - ), + acquire: (_session, _signal, _authorization, _telemetry, use) => + Ref.update(acquireCalls, (count) => count + 1).pipe( + Effect.flatMap(() => use(undefined, Effect.void)), + ), provider: AIProviderHttp.of({ execute: () => Ref.update(providerCalls, (count) => count + 1).pipe( Effect.as({ status: 500, headers: {}, body: null }), @@ -1004,7 +1109,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.fail(AIProviderHttpError.make({ code: "provider_unavailable", @@ -1055,7 +1160,7 @@ describe("Effect AI Gateway forwarding", () => { headers.revoke(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 200, @@ -1105,7 +1210,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 200, @@ -1160,7 +1265,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 200, @@ -1225,7 +1330,7 @@ describe("Effect AI Gateway forwarding", () => { }); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 200, @@ -1254,7 +1359,7 @@ describe("Effect AI Gateway forwarding", () => { const telemetry = yield* makeTelemetryRecorder("upstream_headers"); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 200, @@ -1284,7 +1389,7 @@ describe("Effect AI Gateway forwarding", () => { const settlement = yield* makeSettlementRecorder(); const handler = yield* makeAIForwardHandler({ authentication: { kind: "workload_identity" }, - acquire: () => Effect.succeed(route.lease), + acquire: acquireLease(route.lease), provider: AIProviderHttp.of({ execute: () => Effect.succeed({ status: 200, diff --git a/services/ai-gateway/tests/gateway-service.effect.test.ts b/services/ai-gateway/tests/gateway-service.effect.test.ts index a9ef4b1b..5239b220 100644 --- a/services/ai-gateway/tests/gateway-service.effect.test.ts +++ b/services/ai-gateway/tests/gateway-service.effect.test.ts @@ -10,7 +10,7 @@ import { type ProviderBudgetSettlementReceiptV1, type ProviderBudgetSettlementReportV1, } from "@akua-dev/agentos"; -import { Effect, Ref, Stream } from "effect"; +import { Deferred, Effect, Fiber, Ref, Stream } from "effect"; import { TestClock } from "effect/testing"; import { @@ -117,6 +117,15 @@ const makeTestServices = Effect.fn("test.aiGateway.makeServices")( >>([]); const released = yield* Ref.make(0); const quotaCalls = yield* Ref.make(0); + const reservation = withAccounts + ? { + accountId: "managed-a", + leaseToken: "lease-a", + expiresAt: now + 60_000, + decisionReason: "best_candidate", + } + : undefined; + const releaseReservation = Ref.update(released, (count) => count + 1); const vault = ManagedAccountVault.of({ list: Effect.succeed(withAccounts ? [{ @@ -140,14 +149,22 @@ const makeTestServices = Effect.fn("test.aiGateway.makeServices")( activeReservations: 0, reservationsByAccount: {}, }), - acquire: () => Effect.succeed(withAccounts - ? { - accountId: "managed-a", - leaseToken: "lease-a", - expiresAt: now + 60_000, - decisionReason: "best_candidate", - } - : undefined), + acquire: (_input, use) => Effect.gen(function*() { + const transferred = yield* Ref.make(false); + return yield* Effect.acquireUseRelease( + Effect.succeed(reservation), + (acquired) => use( + acquired, + Effect.uninterruptible(Ref.set(transferred, true)), + ), + (acquired, exit) => Effect.gen(function*() { + if (acquired === undefined) return; + const wasTransferred = yield* Ref.get(transferred); + if (wasTransferred && exit._tag === "Success") return; + yield* releaseReservation; + }), + ); + }), evaluate: () => Effect.succeed(withAccounts ? { accountId: "managed-a", @@ -159,7 +176,7 @@ const makeTestServices = Effect.fn("test.aiGateway.makeServices")( candidates: [], }), renew: () => Effect.succeed(true), - release: () => Ref.update(released, (count) => count + 1).pipe( + release: () => releaseReservation.pipe( Effect.as(true), ), recordResponse: () => Effect.void, @@ -379,6 +396,42 @@ describe("Effect AI Gateway application", () => { }]); })); + it.effect("releases a reservation when credential acquisition is interrupted", () => + Effect.gen(function*() { + yield* TestClock.setTime(now); + const services = yield* makeTestServices(true); + const credentialStarted = yield* Deferred.make(); + const credentialCalls = yield* Ref.make(0); + const application = yield* makeApplication({ + ...services, + vault: ManagedAccountVault.of({ + ...services.vault, + getFreshCredential: () => Effect.gen(function*() { + const call = yield* Ref.modify( + credentialCalls, + (count): readonly [number, number] => [count, count + 1], + ); + if (call === 0) { + return { + providerAccountId: "provider-a", + accessToken: "oauth-provider-secret", + expiresAt: now + 60_000, + }; + } + yield* Deferred.succeed(credentialStarted, undefined); + return yield* Effect.never; + }), + }), + }); + const requestFiber = yield* Effect.forkChild(Effect.exit( + application.handle(providerRequest()), + )); + yield* Deferred.await(credentialStarted); + yield* Fiber.interrupt(requestFiber); + + assert.strictEqual(yield* Ref.get(services.released), 1); + })); + it.effect("preserves serving behavior with native telemetry enabled and disabled", () => Effect.forEach([true, false], (enabled) => Effect.gen(function*() { diff --git a/services/ai-gateway/tests/provider-http.effect.test.ts b/services/ai-gateway/tests/provider-http.effect.test.ts index 7b70a049..5fe1013f 100644 --- a/services/ai-gateway/tests/provider-http.effect.test.ts +++ b/services/ai-gateway/tests/provider-http.effect.test.ts @@ -1,6 +1,16 @@ import { assert, describe, it } from "@effect/vitest"; -import { Effect, Layer, Option, Ref, Stream } from "effect"; import { + Deferred, + Effect, + Exit, + Fiber, + Layer, + Option, + Ref, + Stream, +} from "effect"; +import { + FetchHttpClient, HttpClient, HttpClientError, HttpClientRequest, @@ -11,6 +21,8 @@ import { AIProviderHttp, AIProviderHttpError, AIProviderHttpLive, + AIProviderHttpRequestInit, + makeAIProviderHttpLive, } from "../src/provider-http.ts"; function providerLayer( @@ -96,6 +108,152 @@ describe("AI provider HTTP adapter", () => { }); })); + it.effect("keeps credential-bearing redirects manual at the fetch boundary", () => + Effect.gen(function*() { + const calls: Array<{ + readonly authorization: string | null; + readonly redirect: RequestInit["redirect"] | "default"; + readonly url: string; + }> = []; + const fetchImpl = Object.assign( + (input: string | Request | URL, init?: RequestInit) => { + const url = input instanceof Request ? input.url : input.toString(); + calls.push({ + authorization: new Headers(init?.headers).get("authorization"), + redirect: init?.redirect ?? "default", + url, + }); + if (init?.redirect !== "manual") { + calls.push({ + authorization: new Headers(init?.headers).get("authorization"), + redirect: init?.redirect ?? "default", + url: "https://redirect-target.invalid/v1/responses", + }); + } + return Promise.resolve(new Response(null, { + status: 307, + headers: { location: "https://redirect-target.invalid/v1/responses" }, + })); + }, + { preconnect: globalThis["fetch"].preconnect }, + ); + const fetchLayer = FetchHttpClient.layer.pipe( + Layer.provide(AIProviderHttpRequestInit), + Layer.provide(Layer.succeed(FetchHttpClient.Fetch, fetchImpl)), + ); + const layer = AIProviderHttpLive.pipe( + Layer.provide(fetchLayer), + ); + const response = yield* Effect.gen(function*() { + const provider = yield* AIProviderHttp; + return yield* provider.execute(new Request( + "https://api.openai.test/v1/responses", + { + method: "POST", + headers: { authorization: "Bearer projected-workload-token" }, + body: "{}", + }, + )); + }).pipe(Effect.provide(layer)); + + assert.strictEqual(response.status, 307); + assert.deepStrictEqual(calls, [{ + authorization: "Bearer projected-workload-token", + redirect: "manual", + url: "https://api.openai.test/v1/responses", + }]); + })); + + it.effect("scopes manual redirects to the provider client", () => + Effect.gen(function*() { + const redirects: Array = []; + const fetchImpl = Object.assign( + (_input: string | Request | URL, init?: RequestInit) => { + redirects.push(init?.redirect ?? "default"); + return Promise.resolve(new Response(null, { status: 204 })); + }, + { preconnect: globalThis["fetch"].preconnect }, + ); + const ordinaryClientLayer = FetchHttpClient.layer.pipe( + Layer.provide(Layer.succeed(FetchHttpClient.Fetch, fetchImpl)), + ); + yield* Effect.gen(function*() { + const ordinaryClient = yield* HttpClient.HttpClient; + yield* Effect.scoped( + HttpClient.withScope(ordinaryClient).execute( + HttpClientRequest.get("https://api.openai.test/status"), + ), + ); + }).pipe(Effect.provide(ordinaryClientLayer)); + yield* Effect.gen(function*() { + const provider = yield* AIProviderHttp; + yield* provider.execute(new Request( + "https://api.openai.test/v1/responses", + )); + }).pipe( + Effect.provide(makeAIProviderHttpLive(ordinaryClientLayer)), + ); + + assert.deepStrictEqual(redirects, ["default", "manual"]); + })); + + it.effect("aborts the transport when the caller aborts before upstream headers", () => + Effect.gen(function*() { + const transportSignal = yield* Deferred.make(); + const layer = providerLayer((_request, _url, signal) => + Deferred.succeed(transportSignal, signal).pipe( + Effect.andThen(Effect.never), + )); + const controller = new AbortController(); + const provider = yield* AIProviderHttp.pipe(Effect.provide(layer)); + const fiber = yield* Effect.forkChild(Effect.exit(provider.execute( + new Request("https://api.openai.test/v1/responses", { + method: "POST", + signal: controller.signal, + }), + ))); + const signal = yield* Deferred.await(transportSignal); + yield* Effect.sync(() => controller.abort()); + const exit = yield* Fiber.join(fiber); + + assert.isFalse(Exit.isSuccess(exit)); + assert.isTrue(signal.aborted); + })); + + it.effect("interrupts an in-flight provider stream when the caller aborts", () => + Effect.gen(function*() { + const transportSignal = yield* Deferred.make(); + let bodyStarted = false; + const layer = providerLayer((request, _url, signal) => + Deferred.succeed(transportSignal, signal).pipe( + Effect.andThen(Effect.succeed(HttpClientResponse.fromWeb( + request, + new Response(new ReadableStream({ + start(controller) { + bodyStarted = true; + controller.enqueue(new Uint8Array([1])); + }, + })), + ))), + )); + const controller = new AbortController(); + const provider = yield* AIProviderHttp.pipe(Effect.provide(layer)); + const response = yield* provider.execute(new Request( + "https://api.openai.test/v1/responses", + { signal: controller.signal }, + )); + const bodyFiber = yield* Effect.forkChild(Effect.exit( + Stream.runDrain(response.body ?? Stream.empty), + )); + while (!bodyStarted) yield* Effect.yieldNow; + const signal = yield* Deferred.await(transportSignal); + yield* Effect.sync(() => controller.abort()); + const exit = yield* Fiber.join(bodyFiber); + + assert.isFalse(Exit.isSuccess(exit)); + assert.isTrue(signal.aborted); + })); + it.effect("maps request construction and transport failures to closed typed errors", () => Effect.forEach(transportCases, ({ cause, expected }) => Effect.gen(function*() { diff --git a/services/ai-gateway/tests/routing-state.effect.test.ts b/services/ai-gateway/tests/routing-state.effect.test.ts index dcb5b6b4..07235881 100644 --- a/services/ai-gateway/tests/routing-state.effect.test.ts +++ b/services/ai-gateway/tests/routing-state.effect.test.ts @@ -1,9 +1,21 @@ import * as BunFileSystem from "@effect/platform-bun/BunFileSystem"; +import { + AccountId, + LeaseToken, + RouteLease, + RoutingState as CodexRoutingState, + RoutingStateError, + RoutingSummary, +} from "@akua-dev/codex-router/core"; import { assert, describe, it } from "@effect/vitest"; import { Effect, + Fiber, FileSystem, + Layer, + Option, } from "effect"; +import { TestClock } from "effect/testing"; import { makeEffectAIRoutingStateLayer } from "../src/effect-routing-state.ts"; import { defaultRoutingConfig } from "../src/selection.ts"; @@ -48,11 +60,16 @@ function useRouting( path: string, config: RoutingConfig, operation: (routing: AIRoutingState["Service"]) => Effect.Effect, + routingLayer?: Layer.Layer, ) { return Effect.scoped(Effect.gen(function*() { const routing = yield* AIRoutingState; return yield* operation(routing); - }).pipe(Effect.provide(makeEffectAIRoutingStateLayer(path, config)))); + }).pipe(Effect.provide(makeEffectAIRoutingStateLayer( + path, + config, + routingLayer, + )))); } describe("Effect canonical durable routing state", () => { @@ -68,33 +85,30 @@ describe("Effect canonical durable routing state", () => { defaultRoutingConfig, (routing) => Effect.gen(function*() { - const acquired = yield* routing.acquire({ + return yield* routing.acquire({ candidates, now, sessionKey: "session-1", - }); - assert.isDefined(acquired); - if (acquired === undefined) return undefined; - assert.strictEqual(acquired.accountId, "a"); - assert.strictEqual( - yield* routing.renew(acquired.leaseToken, now + 30_000), - true, - ); - const summary = yield* routing.summary(now + 1); - assert.strictEqual(summary.activeReservations, 1); - assert.deepStrictEqual(summary.reservationsByAccount, { a: 1 }); - assert.strictEqual(summary.lastSelection?.reason, "best_candidate"); - assert.strictEqual( - summary.lastSelection?.candidates.some((candidate) => - candidate.accountId === "a" && candidate.eligible - ), - true, - ); - assert.strictEqual( - yield* routing.release(acquired.leaseToken), - true, - ); - return acquired; + }, (acquired) => Effect.gen(function*() { + assert.isDefined(acquired); + if (acquired === undefined) return undefined; + assert.strictEqual(acquired.accountId, "a"); + assert.strictEqual( + yield* routing.renew(acquired.leaseToken, now + 30_000), + true, + ); + const summary = yield* routing.summary(now + 1); + assert.strictEqual(summary.activeReservations, 1); + assert.deepStrictEqual(summary.reservationsByAccount, { a: 1 }); + assert.strictEqual(summary.lastSelection?.reason, "best_candidate"); + assert.strictEqual( + summary.lastSelection?.candidates.some((candidate) => + candidate.accountId === "a" && candidate.eligible + ), + true, + ); + return acquired; + })); }), ); assert.isDefined(first); @@ -107,11 +121,55 @@ describe("Effect canonical durable routing state", () => { candidates, now: now + 30_001, sessionKey: "session-1", - }), + }, (acquired) => Effect.succeed(acquired)), ); assert.strictEqual(sticky?.accountId, "a"); }).pipe(Effect.provide(BunFileSystem.layer)))); + it.effect("preserves successful use when lease cleanup fails", () => + Effect.gen(function*() { + const fiber = yield* Effect.forkChild(useRouting( + "unused", + defaultRoutingConfig, + (routing) => + routing.acquire({ candidates, now }, (acquired) => + Effect.gen(function*() { + assert.isDefined(acquired); + if (acquired === undefined) { + return new Response(null, { status: 503 }); + } + return new Response(null, { + status: 429, + headers: { "retry-after": "7" }, + }); + }) + ), + Layer.succeed(CodexRoutingState, { + acquire: () => Effect.succeed(Option.some(RouteLease.make({ + accountId: AccountId.make("a"), + expiresAt: now + 60_000, + leaseToken: LeaseToken.make("lease-token"), + sessionKey: Option.none(), + }))), + recordResponse: () => Effect.void, + release: () => Effect.fail(RoutingStateError.make({ + message: "cleanup failed", + })), + renew: () => Effect.succeed(true), + summary: () => Effect.succeed(RoutingSummary.make({ + activeReservations: 0, + accounts: [], + assignments: 0, + })), + }), + )); + yield* Effect.yieldNow; + yield* TestClock.adjust(5_001); + const response = yield* Fiber.join(fiber); + assert.strictEqual(response.status, 429); + assert.strictEqual(response.headers.get("retry-after"), "7"); + })); + it.effect("expires leases without inventing implicit stickiness", () => Effect.scoped(Effect.gen(function*() { const fileSystem = yield* FileSystem.FileSystem; @@ -123,11 +181,19 @@ describe("Effect canonical durable routing state", () => { { ...defaultRoutingConfig, reservationTtlMs: 10 }, (routing) => Effect.gen(function*() { - assert.isDefined(yield* routing.acquire({ candidates, now })); - yield* routing.acquire({ candidates, now: now + 11 }); - assert.strictEqual( - (yield* routing.summary(now + 11)).activeReservations, - 1, + assert.isDefined(yield* routing.acquire( + { candidates, now }, + (acquired) => Effect.succeed(acquired), + )); + yield* routing.acquire( + { candidates, now: now + 11 }, + (acquired) => Effect.gen(function*() { + assert.isDefined(acquired); + assert.strictEqual( + (yield* routing.summary(now + 11)).activeReservations, + 1, + ); + }), ); }), ); @@ -156,15 +222,17 @@ describe("Effect canonical durable routing state", () => { defaultRoutingConfig, (routing) => Effect.gen(function*() { - const first = yield* routing.acquire({ candidates: tied, now }); - const second = yield* routing.acquire({ - candidates: tied, - now: now + 1, - }); - assert.strictEqual(first?.accountId, "a"); - assert.strictEqual(first?.decisionReason, "best_candidate"); - assert.strictEqual(second?.accountId, "b"); - assert.strictEqual(second?.decisionReason, "best_candidate"); + yield* routing.acquire({ candidates: tied, now }, (first) => + routing.acquire({ + candidates: tied, + now: now + 1, + }, (second) => Effect.sync(() => { + assert.strictEqual(first?.accountId, "a"); + assert.strictEqual(first?.decisionReason, "best_candidate"); + assert.strictEqual(second?.accountId, "b"); + assert.strictEqual(second?.decisionReason, "best_candidate"); + })) + ); }), ); }).pipe(Effect.provide(BunFileSystem.layer)))); @@ -187,18 +255,21 @@ describe("Effect canonical durable routing state", () => { needsReauth: true, }], now, - })); - const summary = yield* routing.summary(now); - assert.strictEqual( - summary.lastSelection?.reason, - "no_eligible_accounts", - ); - assert.deepStrictEqual(summary.lastSelection?.candidates, [{ - accountId: "reauth", - eligible: false, - freshness: "unknown", - rejectionCode: "reauthentication_required", - }]); + }, (acquired) => Effect.gen(function*() { + assert.isUndefined(acquired); + const summary = yield* routing.summary(now); + assert.strictEqual( + summary.lastSelection?.reason, + "no_eligible_accounts", + ); + assert.deepStrictEqual(summary.lastSelection?.candidates, [{ + accountId: "reauth", + eligible: false, + freshness: "unknown", + rejectionCode: "reauthentication_required", + }]); + return acquired; + }))); }), ); }).pipe(Effect.provide(BunFileSystem.layer)))); diff --git a/services/ai-gateway/tests/workload-client-proxy.effect.test.ts b/services/ai-gateway/tests/workload-client-proxy.effect.test.ts new file mode 100644 index 00000000..96d560f6 --- /dev/null +++ b/services/ai-gateway/tests/workload-client-proxy.effect.test.ts @@ -0,0 +1,393 @@ +import * as BunServices from "@effect/platform-bun/BunServices"; +import { assert, layer } from "@effect/vitest"; +import { + ConfigProvider, + Effect, + FileSystem, + Ref, + Schema, + Stream, +} from "effect"; +import { TestClock } from "effect/testing"; +import { parse } from "yaml"; + +import { + loadWorkloadClientProxyConfig, + makeWorkloadClientProxyHandler, + workloadClientProxyErrorResponse, + workloadClientProxyReadinessResponse, +} from "../src/workload-client-proxy.ts"; +import { responseFromUpstream } from "../src/workload-client-proxy-main.ts"; + +const suite = layer(BunServices.layer); +const token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJmaXh0dXJlIn0.signature"; +const assignmentId = "20000000-0000-4000-8000-000000000001"; + +const HermesConfiguration = Schema.Struct({ + model: Schema.Struct({ + provider: Schema.String, + default: Schema.String, + }), + providers: Schema.Struct({ + "agentos-gateway": Schema.Struct({ + name: Schema.String, + api: Schema.String, + api_key: Schema.String, + transport: Schema.String, + discover_models: Schema.Boolean, + models: Schema.Array(Schema.String), + }), + }), + agent: Schema.Struct({ api_max_retries: Schema.Number }), + fallback_providers: Schema.Array(Schema.String), + fallback_model: Schema.String, +}); + +function environment(values: Readonly>) { + return ConfigProvider.layer(ConfigProvider.fromEnv({ env: { ...values } })); +} + +suite("Hermes Responses workload client proxy", (it) => { + it.effect("resolves the pinned Hermes named provider fixture to an allowed Responses path", () => + Effect.gen(function*() { + const fileSystem = yield* FileSystem.FileSystem; + const fixtureUrl = new URL( + "../../../docs/integrations/hermes-ai-gateway.config.yaml", + import.meta.url, + ); + const fixture = yield* fileSystem.readFileString(fixtureUrl.pathname); + const parsedFixture = parse(fixture); + const config = yield* Schema.decodeUnknownEffect(HermesConfiguration)( + parsedFixture, + ); + + assert.strictEqual(config.model.provider, "custom:agentos-gateway"); + const provider = config.providers["agentos-gateway"]; + assert.strictEqual(provider.name, "AgentOS Gateway"); + assert.strictEqual(provider.transport, "codex_responses"); + assert.strictEqual(provider.api, "http://127.0.0.1:8790/v1"); + assert.strictEqual( + provider.api_key, + "agentos-workload-identity-placeholder", + ); + assert.isFalse(provider.discover_models); + assert.deepStrictEqual(provider.models, [config.model.default]); + assert.strictEqual(config.agent.api_max_retries, 0); + assert.deepStrictEqual(config.fallback_providers, []); + assert.strictEqual(config.fallback_model, ""); + assert.isFalse(Object.prototype.hasOwnProperty.call(parsedFixture, "auxiliary")); + + const directory = yield* fileSystem.makeTempDirectoryScoped(); + const tokenPath = `${directory}/token`; + yield* fileSystem.writeFileString(tokenPath, token); + const forwarded = yield* Ref.make(null); + const handler = yield* makeWorkloadClientProxyHandler({ + upstreamBaseUrl: new URL("http://agentgateway-openai.agentos.svc.cluster.local:8788"), + tokenPath, + forward: (request) => + Ref.set(forwarded, request).pipe( + Effect.as(new Response(null, { status: 204 })), + ), + }); + + const hermesResponsesUrl = `${ + provider.api.replace(/\/$/, "") + }/responses`; + const response = yield* handler(new Request(hermesResponsesUrl, { + method: "POST", + headers: { authorization: `Bearer ${provider.api_key}` }, + body: JSON.stringify({ model: config.model.default }), + })); + const request = yield* Ref.get(forwarded); + + assert.strictEqual(response.status, 204); + assert.strictEqual( + request?.url, + "http://agentgateway-openai.agentos.svc.cluster.local:8788/v1/responses", + ); + })); + + it.effect("rereads identity, sanitizes authority headers, and injects only trusted assignment", () => + Effect.gen(function*() { + const fileSystem = yield* FileSystem.FileSystem; + const directory = yield* fileSystem.makeTempDirectoryScoped(); + const tokenPath = `${directory}/token`; + yield* fileSystem.writeFileString(tokenPath, token); + const requests = yield* Ref.make>([]); + const handler = yield* makeWorkloadClientProxyHandler({ + upstreamBaseUrl: new URL("http://agentgateway-openai.agentos.svc.cluster.local:8788"), + tokenPath, + assignmentId, + forward: (request) => + Ref.update(requests, (current) => [...current, request]).pipe( + Effect.as(new Response("provider failure", { + status: 429, + headers: { "x-upstream": "preserved" }, + })), + ), + }); + + const first = yield* handler(new Request("http://127.0.0.1:8790/v1/responses", { + method: "POST", + headers: { + authorization: "Bearer hermes-placeholder", + "content-type": "application/json", + "api-key": "caller-provider-secret", + "x-api-key": "caller-provider-secret-2", + "chatgpt-account-id": "caller-account", + baggage: "caller-baggage", + "x-ai-router-token": "caller-router-token", + "x-codex-router-session": "caller-router-session", + "x-ai-gateway-session": "caller-gateway-session", + traceparent: "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "x-agentos-assignment-id": "forged-assignment", + "x-agentos-decision": "forged-decision", + "x-agentos-grant": "forged-grant", + "x-ai-gateway-token": "legacy-shared-token", + }, + body: "{\"model\":\"gpt-exact\"}", + })); + yield* fileSystem.writeFileString(tokenPath, `${token}a`); + const second = yield* handler(new Request("http://127.0.0.1:8790/v1/responses/compact", { + method: "POST", + body: "{}", + })); + + assert.strictEqual(first.status, 429); + assert.strictEqual(first.headers.get("x-upstream"), "preserved"); + assert.strictEqual(yield* Effect.promise(() => first.text()), "provider failure"); + assert.strictEqual(second.status, 429); + const forwarded = yield* Ref.get(requests); + assert.lengthOf(forwarded, 2); + assert.strictEqual(forwarded[0]?.headers.get("authorization"), `Bearer ${token}`); + assert.strictEqual(forwarded[0]?.headers.get("x-agentos-assignment-id"), assignmentId); + assert.isNull(forwarded[0]?.headers.get("x-agentos-decision")); + assert.isNull(forwarded[0]?.headers.get("x-agentos-grant")); + assert.isNull(forwarded[0]?.headers.get("x-ai-gateway-token")); + assert.isNull(forwarded[0]?.headers.get("api-key")); + assert.isNull(forwarded[0]?.headers.get("x-api-key")); + assert.isNull(forwarded[0]?.headers.get("chatgpt-account-id")); + assert.isNull(forwarded[0]?.headers.get("baggage")); + assert.isNull(forwarded[0]?.headers.get("x-ai-router-token")); + assert.isNull(forwarded[0]?.headers.get("x-codex-router-session")); + assert.isNull(forwarded[0]?.headers.get("x-ai-gateway-session")); + assert.strictEqual( + forwarded[0]?.headers.get("traceparent"), + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ); + assert.strictEqual(forwarded[1]?.headers.get("authorization"), `Bearer ${token}a`); + assert.strictEqual( + yield* Effect.promise(() => forwarded[0]?.text() ?? Promise.resolve("")), + "{\"model\":\"gpt-exact\"}", + ); + })); + + it.effect("hard-codes loopback listening and rejects invalid trusted assignment configuration", () => + Effect.gen(function*() { + const config = yield* loadWorkloadClientProxyConfig().pipe( + Effect.provide(environment({ + AI_GATEWAY_URL: "http://agentgateway-openai.agentos.svc.cluster.local:8788", + AI_GATEWAY_WORKLOAD_PROXY_HOST: "0.0.0.0", + AGENTOS_ASSIGNMENT_ID: assignmentId, + })), + ); + assert.strictEqual(config.hostname, "127.0.0.1"); + assert.strictEqual(config.assignmentId, assignmentId); + assert.strictEqual(config.idleTimeoutSeconds, 255); + assert.strictEqual(config.gracefulShutdownMillis, 20_000); + + const configured = yield* loadWorkloadClientProxyConfig().pipe( + Effect.provide(environment({ + AI_GATEWAY_URL: "http://agentgateway-openai.agentos.svc.cluster.local:8788", + AI_GATEWAY_IDLE_TIMEOUT_SECONDS: "120", + AI_GATEWAY_GRACEFUL_SHUTDOWN_MILLIS: "45000", + })), + ); + assert.strictEqual(configured.idleTimeoutSeconds, 120); + assert.strictEqual(configured.gracefulShutdownMillis, 45_000); + + for (const invalidUrl of [ + "ftp://agentgateway-openai.agentos.svc.cluster.local:8788", + "http://agentgateway-github.agentos.svc.cluster.local:8788", + "http://example.invalid:8788", + "http://user:password@agentgateway-openai.agentos.svc.cluster.local:8788", + "http://agentgateway-openai.agentos.svc.cluster.local:8788?trace=1", + "http://agentgateway-openai.agentos.svc.cluster.local:8788#fragment", + "http://agentgateway-openai.agentos.svc.cluster.local:8788/v1", + ]) { + const invalid = yield* loadWorkloadClientProxyConfig().pipe( + Effect.provide(environment({ AI_GATEWAY_URL: invalidUrl })), + Effect.flip, + ); + assert.strictEqual(invalid.code, "invalid_configuration"); + } + + for (const invalidAssignmentId of [ + "caller-controlled-value", + "20000000-0000-1000-8000-000000000001", + "20000000-0000-4000-7000-000000000001", + ]) { + const invalid = yield* loadWorkloadClientProxyConfig().pipe( + Effect.provide(environment({ + AI_GATEWAY_URL: "http://agentgateway-openai.agentos.svc.cluster.local:8788", + AGENTOS_ASSIGNMENT_ID: invalidAssignmentId, + })), + Effect.flip, + ); + assert.strictEqual(invalid.code, "invalid_configuration"); + } + })); + + it.effect("strips hop-by-hop response headers while preserving provider data", () => + Effect.gen(function*() { + const response = yield* responseFromUpstream({ + status: 429, + headers: { + connection: "keep-alive, x-connection-hop", + "keep-alive": "timeout=5", + "proxy-authenticate": "Basic", + "proxy-authorization": "Basic secret", + te: "trailers", + trailer: "x-trailer", + "transfer-encoding": "chunked", + upgrade: "websocket", + "x-connection-hop": "private", + "x-provider-request-id": "provider-1", + }, + body: Stream.succeed(new TextEncoder().encode("provider-body")), + }); + + assert.strictEqual(response.status, 429); + for (const name of [ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + "x-connection-hop", + ]) { + assert.isNull(response.headers.get(name)); + } + assert.strictEqual( + response.headers.get("x-provider-request-id"), + "provider-1", + ); + assert.strictEqual( + yield* Effect.promise(() => response.text()), + "provider-body", + ); + })); + + it.effect("preserves a 205 response without constructing a body", () => + Effect.gen(function*() { + const response = yield* responseFromUpstream({ + status: 205, + headers: { "content-type": "text/plain" }, + body: Stream.succeed(new TextEncoder().encode("invalid-body")), + }); + + assert.strictEqual(response.status, 205); + assert.isNull(response.body); + assert.strictEqual( + yield* Effect.promise(() => response.text()), + "", + ); + })); + + it.effect("preserves upstream response streams across a long gap", () => + Effect.gen(function*() { + const fileSystem = yield* FileSystem.FileSystem; + const directory = yield* fileSystem.makeTempDirectoryScoped(); + const tokenPath = `${directory}/token`; + yield* fileSystem.writeFileString(tokenPath, token); + const handler = yield* makeWorkloadClientProxyHandler({ + upstreamBaseUrl: new URL("http://agentgateway-openai.agentos.svc.cluster.local:8788"), + tokenPath, + forward: () => Stream.toReadableStreamEffect(Stream.concat( + Stream.succeed(new TextEncoder().encode("first")), + Stream.fromEffect( + Effect.sleep("1.1 seconds").pipe( + Effect.andThen(Effect.succeed(new TextEncoder().encode("second"))), + ), + ), + )).pipe( + Effect.map((body) => new Response(body, { + status: 200, + headers: { "content-type": "text/plain" }, + })), + ), + }); + + const response = yield* TestClock.withLive(handler(new Request( + "http://127.0.0.1:8790/v1/responses", + { method: "POST" }, + ))); + assert.strictEqual(response.status, 200); + assert.strictEqual( + yield* TestClock.withLive(Effect.promise(() => response.text())), + "firstsecond", + ); + })); + + it.effect("reports readiness only for a bounded, trimmed, UTF-8 JWT-like token", () => + Effect.gen(function*() { + const fileSystem = yield* FileSystem.FileSystem; + const directory = yield* fileSystem.makeTempDirectoryScoped(); + const tokenPath = `${directory}/token`; + + const missing = yield* workloadClientProxyReadinessResponse(tokenPath); + assert.strictEqual(missing.status, 503); + + const invalidTokens: ReadonlyArray = [ + new Uint8Array(), + new TextEncoder().encode(`${token}\n`), + new TextEncoder().encode("not-a-jwt"), + new Uint8Array([0xff, 0xfe, 0xfd]), + new Uint8Array(16 * 1024 + 1).fill(97), + ]; + for (const bytes of invalidTokens) { + yield* fileSystem.writeFile(tokenPath, bytes); + const response = yield* workloadClientProxyReadinessResponse(tokenPath); + assert.strictEqual(response.status, 503); + } + + yield* fileSystem.writeFileString(tokenPath, token); + const ready = yield* workloadClientProxyReadinessResponse(tokenPath); + assert.strictEqual(ready.status, 200); + assert.deepStrictEqual(yield* Effect.promise(() => ready.json()), { + status: "ready", + }); + })); + + it.effect("fails closed without sending unsupported, unauthenticated, or readiness requests", () => + Effect.gen(function*() { + const fileSystem = yield* FileSystem.FileSystem; + const directory = yield* fileSystem.makeTempDirectoryScoped(); + const tokenPath = `${directory}/token`; + const calls = yield* Ref.make(0); + const handler = yield* makeWorkloadClientProxyHandler({ + upstreamBaseUrl: new URL("http://agentgateway-openai.agentos.svc.cluster.local:8788"), + tokenPath, + forward: () => Ref.update(calls, (value) => value + 1).pipe( + Effect.as(new Response(null, { status: 200 })), + ), + }); + + const unsupported = yield* handler(new Request("http://127.0.0.1:8790/v1/chat/completions", { + method: "POST", + })).pipe(Effect.flip); + assert.strictEqual(unsupported.code, "invalid_request"); + assert.strictEqual(workloadClientProxyErrorResponse(unsupported).status, 404); + + const missing = yield* handler(new Request("http://127.0.0.1:8790/v1/responses", { + method: "POST", + })).pipe(Effect.flip); + assert.strictEqual(missing.code, "token_unavailable"); + assert.strictEqual(workloadClientProxyErrorResponse(missing).status, 503); + yield* workloadClientProxyReadinessResponse(tokenPath); + assert.strictEqual(yield* Ref.get(calls), 0); + })); +}); diff --git a/tooling/effect-migration/exceptions.json b/tooling/effect-migration/exceptions.json index 8acac5e4..a248fc6c 100644 --- a/tooling/effect-migration/exceptions.json +++ b/tooling/effect-migration/exceptions.json @@ -433,6 +433,18 @@ "removalCondition": "Retain only while the AI Gateway remains a standalone Bun service and keep all domain behavior inside Effect.", "expiresOn": null }, + { + "kind": "outer-host-adapter", + "path": "services/ai-gateway/src/workload-client-proxy-main.ts", + "rule": "no-runtime-execution", + "match": "BunRuntime.runMain", + "maximumOccurrences": 1, + "reason": "The workload client proxy executable needs one reviewed Bun process entry boundary around its fully typed, scoped Effect Platform application.", + "ownerIssue": 105, + "test": "services/ai-gateway/tests/workload-client-proxy.effect.test.ts", + "removalCondition": "Retain only while the loopback workload client proxy remains a standalone Bun sidecar entrypoint and keep all request behavior inside Effect.", + "expiresOn": null + }, { "kind": "outer-host-adapter", "path": "services/egress-authz/src/main.ts", diff --git a/tooling/effect-migration/inventory.json b/tooling/effect-migration/inventory.json index 40864338..eff4b5f4 100644 --- a/tooling/effect-migration/inventory.json +++ b/tooling/effect-migration/inventory.json @@ -963,7 +963,7 @@ { "id": "ai-gateway-entrypoint", "issue": 105, - "pattern": "^services/ai-gateway/src/main\\.ts$", + "pattern": "^services/ai-gateway/src/(?:main|workload-client-proxy-main)\\.ts$", "runtime": "bun-entrypoint", "package": "@agentos/ai-gateway", "io": ["config", "filesystem", "http", "oauth", "process", "stream", "telemetry", "runtime-adapter"], diff --git a/website/apps/docs/content/docs/operate/ai-gateway.mdx b/website/apps/docs/content/docs/operate/ai-gateway.mdx index a67eb515..561b2e2c 100644 --- a/website/apps/docs/content/docs/operate/ai-gateway.mdx +++ b/website/apps/docs/content/docs/operate/ai-gateway.mdx @@ -6,6 +6,8 @@ canonical: path: packages/agentos/skills/agentos-ai-gateway/SKILL.md - label: AI Gateway service path: services/ai-gateway + - label: Hermes integration contract + path: docs/integrations/hermes-ai-gateway.md --- The Fleet AI Gateway is optional pooled capacity for reviewed Codex subscription accounts with an @@ -17,13 +19,15 @@ retiring the gateway. It keeps server-owned provider OAuth and quota-aware routi deliberate service boundary. Client Agents receive only the scoped connection material intended for them; they never receive a provider credential or Fleet-shared inference token. -Clients may live in `agentos` or a Second-Mate domain namespace. They call +Native AgentOS clients may live in `agentos` or a Second-Mate domain namespace. They call `agentgateway-openai.agentos.svc.cluster.local` and authenticate with their kubelet-rotated, -audience-bound ServiceAccount token. Agentgateway asks the TokenReview/OpenFGA authorizer for a -closed grant lasting at most 15 seconds. Only Agentgateway can reach the AI Gateway through its -NetworkPolicy; provider OAuth, the operator Secret and the Gateway vault remain in the core -namespace. This lets a Second Mate manage Crewmates in its namespace without gaining another -domain’s identity or provider credentials, while unrelated Internet traffic stays direct. +audience-bound ServiceAccount token. A Responses client that cannot reread that token for every +request uses the loopback-only workload proxy documented in the Hermes integration contract. +Agentgateway asks the TokenReview/OpenFGA authorizer for a closed grant lasting at most 15 seconds. +Only Agentgateway can reach the AI Gateway through its NetworkPolicy; provider OAuth, the operator +Secret and the Gateway vault remain in the core namespace. This lets a Second Mate manage +Crewmates in its namespace without gaining another domain’s identity or provider credentials, +while unrelated Internet traffic stays direct. ## What it does not own