fix(ai): support custom image auth

packages/ai/src/route/auth.ts

@@ -1,6 +1,6 @@
 import { Config, Effect, Redacted } from "effect"
 import { Headers } from "effect/unstable/http"
-import { AuthenticationReason, InvalidRequestReason, LLMError, LLMRequest, type HttpOptions } from "../schema"
+import { AuthenticationReason, InvalidRequestReason, LLMError, type HttpOptions, type LLMRequest } from "../schema"
 
 export class MissingCredentialError extends Error {
   readonly _tag = "MissingCredentialError"
@@ -14,7 +14,7 @@ export type CredentialError = MissingCredentialError | Config.ConfigError
 export type AuthError = CredentialError | LLMError
 type Secret = string | Redacted.Redacted | Config.Config<string | Redacted.Redacted>
 
-interface RequestInput {
+export interface RequestAuthInput {
   readonly request: { readonly http?: HttpOptions }
   readonly method: "POST" | "GET"
   readonly url: string
@@ -22,7 +22,7 @@ interface RequestInput {
   readonly headers: Headers.Headers
 }
 
-export interface AuthInput extends Omit<RequestInput, "request"> {
+export interface AuthInput extends Omit<RequestAuthInput, "request"> {
   readonly request: LLMRequest
 }
 
@@ -35,7 +35,7 @@ export interface Credential {
 }
 
 export interface Definition {
-  readonly apply: (input: RequestInput) => Effect.Effect<Headers.Headers, AuthError>
+  readonly apply: (input: RequestAuthInput) => Effect.Effect<Headers.Headers, AuthError>
   readonly andThen: (that: Definition) => Definition
   readonly orElse: (that: Definition) => Definition
   readonly pipe: <A>(f: (self: Definition) => A) => A
@@ -104,17 +104,11 @@ export const headers = (input: Headers.Input) =>
 
 export const remove = (name: string) => auth((input) => Effect.succeed(Headers.remove(input.headers, name)))
 
+export const customRequest = (apply: (input: RequestAuthInput) => Effect.Effect<Headers.Headers, LLMError>) =>
+  auth(apply)
+
 export const custom = (apply: (input: AuthInput) => Effect.Effect<Headers.Headers, LLMError>) =>
-  auth((input) => {
-    if (input.request instanceof LLMRequest) return apply({ ...input, request: input.request })
-    return Effect.fail(
-      new LLMError({
-        module: "Auth",
-        method: "custom",
-        reason: new InvalidRequestReason({ message: "Custom LLM auth requires an LLM request" }),
-      }),
-    )
-  })
+  customRequest((input) => apply(input as AuthInput))
 
 export const passthrough = none
 
@@ -164,7 +158,7 @@ const toLLMError = (error: AuthError): LLMError => {
 
 export const toEffect =
   (input: Definition) =>
-  (authInput: RequestInput): Effect.Effect<Headers.Headers, LLMError> =>
+  (authInput: RequestAuthInput): Effect.Effect<Headers.Headers, LLMError> =>
     input.apply(authInput).pipe(Effect.mapError(toLLMError))
 
 export * as Auth from "./auth"

packages/ai/src/route/index.ts

@@ -17,7 +17,7 @@ export { Framing } from "./framing"
 export { Protocol } from "./protocol"
 export { HttpTransport, WebSocketExecutor, WebSocketTransport } from "./transport"
 export * as Transport from "./transport"
-export type { Definition as AuthShape, AuthInput, Credential, CredentialError } from "./auth"
+export type { Definition as AuthShape, AuthInput, RequestAuthInput, Credential, CredentialError } from "./auth"
 export type { ApiKeyMode, AuthOverride, ProviderAuthOption } from "./auth-options"
 export type { Definition as EndpointFn, EndpointInput } from "./endpoint"
 export type { Definition as FramingDef } from "./framing"

packages/ai/test/provider/xai-images.test.ts

@@ -1,8 +1,9 @@
 import { describe, expect } from "bun:test"
 import { Effect, Layer } from "effect"
-import { HttpClientRequest } from "effect/unstable/http"
+import { Headers, HttpClientRequest } from "effect/unstable/http"
 import { Image, ImageClient } from "../../src"
 import { XAI } from "../../src/providers"
+import { Auth } from "../../src/route"
 import { it } from "../lib/effect"
 import { dynamicResponse } from "../lib/http"
 
@@ -61,4 +62,32 @@ describe("xAI Images", () => {
       ),
     ),
   )
+
+  it.effect("supports request-level custom auth", () =>
+    Image.generate({
+      model: XAI.configure({
+        baseURL: "https://api.xai.test/v1",
+        auth: Auth.customRequest((input) =>
+          Effect.succeed(Headers.set(input.headers, "x-custom-auth", new URL(input.url).hostname)),
+        ),
+      }).image("grok-imagine-image"),
+      prompt: "A robot tending a rooftop garden",
+    }).pipe(
+      Effect.provide(
+        ImageClient.layer.pipe(
+          Layer.provide(
+            dynamicResponse((input) =>
+              Effect.gen(function* () {
+                const request = yield* HttpClientRequest.toWeb(input.request).pipe(Effect.orDie)
+                expect(request.headers.get("x-custom-auth")).toBe("api.xai.test")
+                return input.respond(JSON.stringify({ data: [{ b64_json: "AQID", mime_type: "image/png" }] }), {
+                  headers: { "content-type": "application/json" },
+                })
+              }),
+            ),
+          ),
+        ),
+      ),
+    ),
+  )
 })