Workshop: loop gateway

package dev.loop.gateway;

import jakarta.validation.Valid;
import jakarta.validation.constraints.Min;
import jakarta.validation.constraints.NotBlank;
import jakarta.validation.constraints.NotEmpty;
import java.time.Duration;
import java.util.List;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.context.properties.bind.DefaultValue;
import org.springframework.validation.annotation.Validated;

@Validated
@ConfigurationProperties("loop.gateway")
public record GatewayProperties(
        @DefaultValue List<String> apiKeys,
        @DefaultValue("5m") Duration idleTimeout,
        @DefaultValue("60s") Duration startTimeout,
        @DefaultValue("1") @Min(1) int maxLoaded,
        @NotEmpty List<@Valid ModelSpec> models) {

    /** One model the gateway serves: a remote server (url) or a command to start (command, with {port}). */
    public record ModelSpec(@NotBlank String name, String url, List<String> command) {
        public ModelSpec {
            if ((url == null) == (command == null || command.isEmpty())) {
                throw new IllegalArgumentException("model " + name + ": set exactly one of url and command");
            }
        }
    }
}

package dev.loop.gateway;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.net.ConnectException;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.ArrayDeque;
import java.util.List;

/** An engine the gateway forwards to. Same shape as the Backend of Project 12. */
public interface Backend extends AutoCloseable {
    String baseUrl();                       // ends with /v1
    void ensureReady() throws Exception;    // start if needed, then wait until healthy
    @Override void close();                 // stop only what this backend started

    /** A server that already runs. */
    record RemoteBackend(String baseUrl) implements Backend {
        public void ensureReady() { }
        public void close() { }
    }

    /** An engine started as a child process on a free port; the command contains {port}. */
    final class ProcessBackend implements Backend {
        private final String name;
        private final List<String> command;
        private final Duration startTimeout;
        private final int port;
        private final ArrayDeque<String> log = new ArrayDeque<>();
        private Process process;

        public ProcessBackend(String name, List<String> command, Duration startTimeout) {
            this.name = name;
            this.startTimeout = startTimeout;
            this.port = freePort();
            this.command = command.stream().map(a -> a.replace("{port}", Integer.toString(port))).toList();
        }

        public String baseUrl() { return "http://127.0.0.1:" + port + "/v1"; }

        public void ensureReady() throws Exception {
            ProcessBuilder pb = new ProcessBuilder(command).redirectErrorStream(true);
            pb.environment().remove("JAVA_TOOL_OPTIONS");
            process = pb.start();
            Thread.startVirtualThread(() -> {
                try { process.inputReader().lines().forEach(this::remember); }
                catch (UncheckedIOException e) { /* the pipe closed because the engine stopped */ }
            });
            HttpClient http = HttpClient.newHttpClient();
            HttpRequest health = HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port + "/health"))
                    .timeout(Duration.ofSeconds(5)).build();
            long deadline = System.nanoTime() + startTimeout.toNanos();
            long sleep = 50;
            while (true) {
                if (!process.isAlive()) {
                    throw new IllegalStateException("engine " + name + " exited with code " + process.exitValue()
                            + " during startup: " + lastLog());
                }
                try {
                    if (http.send(health, HttpResponse.BodyHandlers.discarding()).statusCode() == 200) return;
                } catch (ConnectException e) {
                    // not listening yet
                }
                if (System.nanoTime() > deadline) {
                    close();
                    throw new IllegalStateException("engine " + name + " not ready after " + startTimeout);
                }
                Thread.sleep(sleep);
                sleep = Math.min(sleep * 2, 1000);
            }
        }

        public void close() {
            if (process == null) return;
            process.descendants().forEach(ProcessHandle::destroy);
            process.destroy();
            try { process.waitFor(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); }
        }

        private synchronized void remember(String line) { if (log.size() == 20) log.removeFirst(); log.addLast(line); }
        private synchronized String lastLog() { return log.isEmpty() ? "(no output)" : log.peekLast(); }

        private static int freePort() {
            try (ServerSocket s = new ServerSocket(0, 0, InetAddress.getLoopbackAddress())) { return s.getLocalPort(); }
            catch (IOException e) { throw new IllegalStateException(e); }
        }
    }
}

package dev.loop.gateway;

import dev.loop.gateway.GatewayProperties.ModelSpec;
import io.micrometer.core.instrument.Gauge;
import io.micrometer.core.instrument.MeterRegistry;
import jakarta.annotation.PreDestroy;
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

@Component
public class BackendManager {
    private static final Logger log = LoggerFactory.getLogger(BackendManager.class);

    /** What a request holds while it uses an engine; closing it releases the engine. */
    public record Lease(String model, String baseUrl, Runnable onClose) implements AutoCloseable {
        @Override public void close() { onClose.run(); }
    }

    public record LoadedModel(String name, int active, Instant lastUsed) {}

    private static final class Entry {
        final CompletableFuture<Backend> ready = new CompletableFuture<>();
        int active;
        Instant lastUsed;
    }

    private final GatewayProperties props;
    private final Clock clock;
    private final Map<String, Entry> entries = new LinkedHashMap<>();

    public BackendManager(GatewayProperties props, Clock clock, MeterRegistry metrics) {
        this.props = props;
        this.clock = clock;
        Gauge.builder("loop.gateway.models.loaded", this, m -> m.loaded().size()).register(metrics);
    }

    public List<ModelSpec> configured() { return props.models(); }

    public Lease acquire(String model) {
        Entry entry;
        boolean starter = false;
        ModelSpec spec = props.models().stream().filter(m -> m.name().equals(model)).findFirst()
                .orElseThrow(() -> new GatewayException(404, "model_not_found", "The model '" + model + "' does not exist"));
        synchronized (this) {
            entry = entries.get(model);
            if (entry == null) {
                makeRoom();
                entry = new Entry();
                entries.put(model, entry);
                starter = true;
            }
            entry.active++;
            entry.lastUsed = clock.instant();
        }
        if (starter) load(spec, entry);
        try {
            Backend backend = entry.ready.join();
            Entry held = entry;
            return new Lease(model, backend.baseUrl(), () -> release(held));
        } catch (CompletionException e) {
            release(entry);
            throw e.getCause() instanceof GatewayException g ? g
                    : new GatewayException(503, "engine_unavailable", e.getCause().getMessage());
        }
    }

    private void load(ModelSpec spec, Entry entry) {
        log.info("loading model {}", spec.name());
        Backend backend = spec.url() != null ? new Backend.RemoteBackend(spec.url())
                : new Backend.ProcessBackend(spec.name(), spec.command(), props.startTimeout());
        try {
            backend.ensureReady();
            entry.ready.complete(backend);
        } catch (Exception e) {
            backend.close();
            synchronized (this) { entries.remove(spec.name(), entry); }
            entry.ready.completeExceptionally(e);
        }
    }

    private synchronized void release(Entry entry) {
        entry.active--;
        entry.lastUsed = clock.instant();
    }

    /** Called with the lock held: free a slot by stopping the least recently used idle model. */
    private void makeRoom() {
        if (entries.size() < props.maxLoaded()) return;
        Map.Entry<String, Entry> victim = entries.entrySet().stream()
                .filter(e -> e.getValue().active == 0 && e.getValue().ready.isDone())
                .min(Comparator.comparing(e -> e.getValue().lastUsed)).orElseThrow(() ->
                        new GatewayException(503, "models_busy", "All " + props.maxLoaded() + " model slots are in use"));
        log.info("evicting model {}", victim.getKey());
        entries.remove(victim.getKey());
        victim.getValue().ready.join().close();
    }

    @Scheduled(fixedDelayString = "${loop.gateway.tick:5s}")
    public void unloadIdle() { unloadIdle(clock.instant()); }

    public synchronized void unloadIdle(Instant now) {
        for (var it = entries.entrySet().iterator(); it.hasNext(); ) {
            var e = it.next();
            Entry entry = e.getValue();
            if (entry.ready.isDone() && entry.active == 0
                    && !entry.lastUsed.plus(props.idleTimeout()).isAfter(now)) {
                log.info("unloading idle model {}", e.getKey());
                it.remove();
                entry.ready.join().close();
            }
        }
    }

    public synchronized List<LoadedModel> loaded() {
        List<LoadedModel> out = new ArrayList<>();
        entries.forEach((name, e) -> { if (e.ready.isDone()) out.add(new LoadedModel(name, e.active, e.lastUsed)); });
        return out;
    }

    public Duration idleTimeout() { return props.idleTimeout(); }

    @PreDestroy
    public synchronized void shutdown() {
        entries.values().forEach(e -> { if (e.ready.isDone() && !e.ready.isCompletedExceptionally()) e.ready.join().close(); });
        entries.clear();
    }
}

package dev.loop.gateway;

import java.time.Clock;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.security.autoconfigure.UserDetailsServiceAutoConfiguration;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.scheduling.annotation.EnableScheduling;

@SpringBootApplication(exclude = UserDetailsServiceAutoConfiguration.class)
@EnableConfigurationProperties(GatewayProperties.class)
@EnableScheduling
public class GatewayApp {
    @Bean Clock clock() { return Clock.systemUTC(); }
}

package dev.loop.gateway;

import java.time.Instant;
import java.util.List;
import java.util.Map;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

@RestController
public class ModelsController {
    private final BackendManager manager;

    public ModelsController(BackendManager manager) { this.manager = manager; }

    record ModelObject(String id, String object, long created, String owned_by) {}
    record ModelList(String object, List<ModelObject> data) {}

    @GetMapping("/v1/models")
    ModelList models() {
        return new ModelList("list", manager.configured().stream()
                .map(m -> new ModelObject(m.name(), "model", 0, "loop")).toList());
    }

    @GetMapping("/api/tags")
    Map<String, Object> tags() {
        return Map.of("models", manager.configured().stream()
                .map(m -> Map.of("name", m.name(), "model", m.name())).toList());
    }

    @GetMapping("/api/ps")
    Map<String, Object> ps() {
        return Map.of("models", manager.loaded().stream()
                .map(m -> Map.of("name", m.name(), "model", m.name(),
                        "expires_at", m.lastUsed().plus(manager.idleTimeout()).toString())).toList());
    }
}

package dev.loop.gateway;

import io.micrometer.core.instrument.MeterRegistry;
import java.io.InputStream;
import java.net.ConnectException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.net.http.HttpTimeoutException;
import java.time.Duration;
import org.springframework.http.CacheControl;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
import tools.jackson.databind.JsonNode;
import tools.jackson.databind.json.JsonMapper;

@RestController
public class ChatController {
    private final BackendManager manager;
    private final MeterRegistry metrics;
    private final JsonMapper json;
    private final HttpClient http = HttpClient.newBuilder()
            .version(HttpClient.Version.HTTP_1_1).connectTimeout(Duration.ofSeconds(5)).build();

    public ChatController(BackendManager manager, MeterRegistry metrics, JsonMapper json) {
        this.manager = manager;
        this.metrics = metrics;
        this.json = json;
    }

    @PostMapping("/v1/chat/completions")
    ResponseEntity<StreamingResponseBody> chat(@RequestBody byte[] raw) throws Exception {
        JsonNode body = json.readTree(raw);
        String model = body.path("model").asString("");
        if (model.isEmpty()) throw new GatewayException(400, "invalid_request_error", "The field 'model' is required");
        if (!body.path("messages").isArray() || body.path("messages").isEmpty()) {
            throw new GatewayException(400, "invalid_request_error", "The field 'messages' must be a non-empty array");
        }
        boolean stream = body.path("stream").asBoolean(false);

        BackendManager.Lease lease = manager.acquire(model);          // unknown models fail here
        metrics.counter("loop.gateway.requests", "model", model, "stream", Boolean.toString(stream)).increment();
        boolean handedOver = false;
        try {
            HttpRequest upstream = HttpRequest.newBuilder(URI.create(lease.baseUrl() + "/chat/completions"))
                    .timeout(Duration.ofMinutes(10)).header("Content-Type", "application/json")
                    .POST(HttpRequest.BodyPublishers.ofByteArray(raw)).build();   // the body goes through unchanged
            HttpResponse<InputStream> res = http.send(upstream, HttpResponse.BodyHandlers.ofInputStream());
            boolean sse = res.statusCode() == 200 && stream;
            if (!sse) {
                try (InputStream in = res.body()) {
                    byte[] answer = in.readAllBytes();
                    return ResponseEntity.status(res.statusCode()).contentType(MediaType.APPLICATION_JSON)
                            .body(out -> out.write(answer));
                }
            }
            StreamingResponseBody pump = out -> {
                try (lease; InputStream in = res.body()) {      // closing the stream cancels the upstream request
                    byte[] buf = new byte[8192];
                    for (int n = in.read(buf); n != -1; n = in.read(buf)) {
                        out.write(buf, 0, n);
                        out.flush();
                    }
                }
            };
            handedOver = true;
            return ResponseEntity.ok().contentType(MediaType.TEXT_EVENT_STREAM)
                    .cacheControl(CacheControl.noCache()).header("X-Accel-Buffering", "no").body(pump);
        } catch (ConnectException e) {
            throw new GatewayException(502, "engine_unreachable", "The engine for '" + model + "' refused the connection");
        } catch (HttpTimeoutException e) {
            throw new GatewayException(504, "engine_timeout", "The engine for '" + model + "' did not answer in time");
        } finally {
            if (!handedOver) lease.close();
        }
    }
}

package dev.loop.gateway;

/** An error the gateway reports in the OpenAI error shape. */
public class GatewayException extends RuntimeException {
    private final int status;
    private final String type;

    public GatewayException(int status, String type, String message) {
        super(message);
        this.status = status;
        this.type = type;
    }
    public int status() { return status; }
    public String type() { return type; }
}

package dev.loop.gateway;

import java.util.Map;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.http.converter.HttpMessageNotReadableException;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.RestControllerAdvice;

@RestControllerAdvice
public class ErrorAdvice {
    static Map<String, Object> body(String type, String message) {
        return Map.of("error", Map.of("message", message, "type", type));
    }

    @ExceptionHandler(GatewayException.class)
    ResponseEntity<Map<String, Object>> gateway(GatewayException e) {
        return ResponseEntity.status(e.status()).body(body(e.type(), e.getMessage()));
    }

    @ExceptionHandler({HttpMessageNotReadableException.class, tools.jackson.core.JacksonException.class})
    ResponseEntity<Map<String, Object>> unreadable(Exception e) {
        return ResponseEntity.status(HttpStatus.BAD_REQUEST).body(body("invalid_request_error", "The request body is not valid JSON"));
    }
}

package dev.loop.gateway;

import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.List;
import org.springframework.http.MediaType;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.filter.OncePerRequestFilter;

/** Accepts "Authorization: Bearer <key>" when the key is configured; an empty key list turns the check off. */
public class ApiKeyFilter extends OncePerRequestFilter {
    private final List<byte[]> keys;

    public ApiKeyFilter(List<String> keys) {
        this.keys = keys.stream().map(k -> k.getBytes(StandardCharsets.UTF_8)).toList();
    }

    @Override
    protected void doFilterInternal(HttpServletRequest req, HttpServletResponse res, FilterChain chain)
            throws ServletException, IOException {
        String header = req.getHeader("Authorization");
        boolean ok = keys.isEmpty() || (header != null && header.startsWith("Bearer ") && matches(header.substring(7)));
        if (!ok) {
            res.setStatus(401);
            res.setContentType(MediaType.APPLICATION_JSON_VALUE);
            res.setHeader("WWW-Authenticate", "Bearer");
            res.getWriter().write("{\"error\":{\"message\":\"Incorrect API key provided\",\"type\":\"invalid_api_key\"}}");
            return;
        }
        SecurityContextHolder.getContext().setAuthentication(UsernamePasswordAuthenticationToken.authenticated(
                "client", null, List.of(new SimpleGrantedAuthority("ROLE_API"))));
        chain.doFilter(req, res);
    }

    @Override
    protected boolean shouldNotFilter(HttpServletRequest req) {
        return req.getRequestURI().startsWith("/actuator/health");     // the chain permits this path
    }

    private boolean matches(String presented) {
        byte[] p = presented.getBytes(StandardCharsets.UTF_8);
        boolean found = false;
        for (byte[] k : keys) found |= MessageDigest.isEqual(k, p);     // constant time per key
        return found;
    }
}

package dev.loop.gateway;

import jakarta.servlet.DispatcherType;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.http.SessionCreationPolicy;
import org.springframework.security.web.SecurityFilterChain;
import org.springframework.security.web.authentication.UsernamePasswordAuthenticationFilter;

@Configuration
public class SecurityConfig {
    @Bean
    SecurityFilterChain chain(HttpSecurity http, GatewayProperties props) throws Exception {
        return http
                .csrf(c -> c.disable())
                .sessionManagement(s -> s.sessionCreationPolicy(SessionCreationPolicy.STATELESS))
                .authorizeHttpRequests(a -> a
                        .dispatcherTypeMatchers(DispatcherType.ASYNC, DispatcherType.ERROR).permitAll()
                        .requestMatchers("/actuator/health", "/actuator/health/**").permitAll()
                        .anyRequest().authenticated())
                .addFilterBefore(new ApiKeyFilter(props.apiKeys()), UsernamePasswordAuthenticationFilter.class)
                .build();
    }
}

package dev.loop.gateway;

import java.util.List;
import org.springframework.boot.health.contributor.Health;
import org.springframework.boot.health.contributor.HealthIndicator;
import org.springframework.stereotype.Component;

/** UP while the gateway runs; the details list the engines that are loaded now. */
@Component("engines")
public class EnginesHealth implements HealthIndicator {
    private final BackendManager manager;

    public EnginesHealth(BackendManager manager) { this.manager = manager; }

    @Override
    public Health health() {
        List<String> loaded = manager.loaded().stream().map(BackendManager.LoadedModel::name).toList();
        return Health.up().withDetail("configured", manager.configured().size()).withDetail("loaded", loaded).build();
    }
}

package dev.loop.gateway;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import org.springframework.boot.SpringApplication;
import picocli.CommandLine;
import picocli.CommandLine.Command;
import picocli.CommandLine.Option;

@Command(name = "gateway", mixinStandardHelpOptions = true, version = "loop gateway 0.1.0",
        description = "Serve the OpenAI-compatible and Ollama-style API in front of the configured engines.")
public class GatewayCommand implements Callable<Integer> {
    @Option(names = "--config", paramLabel = "FILE", description = "YAML file with the loop.gateway settings.")
    String config;

    @Option(names = {"-p", "--port"}, defaultValue = "11435", description = "Port to listen on (default: ${DEFAULT-VALUE}).")
    int port;

    @Option(names = "--host", defaultValue = "127.0.0.1", description = "Address to bind (default: ${DEFAULT-VALUE}).")
    String host;

    @Override
    public Integer call() {
        List<String> args = new ArrayList<>(List.of("--server.port=" + port, "--server.address=" + host));
        if (config != null) args.add("--spring.config.additional-location=optional:file:" + config);
        SpringApplication.run(GatewayApp.class, args.toArray(String[]::new));
        return 0;
    }

    public static void main(String[] args) {
        int code = new CommandLine(new GatewayCommand()).execute(args);
        if (code != 0) System.exit(code);       // on success the server keeps the JVM alive
    }
}

package dev.loop.gateway;

import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpServer;
import java.io.IOException;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/** A fake OpenAI-compatible engine: java -cp ... dev.loop.gateway.FakeEngine PORT NAME [broken]. */
public class FakeEngine {
    static final Pattern CONTENT = Pattern.compile("\"content\"\\s*:\\s*\"((?:[^\"\\\\]|\\\\.)*)\"");

    public static void main(String[] args) throws Exception {
        int port = Integer.parseInt(args[0]);
        String name = args[1];
        if (args.length > 2 && args[2].equals("broken")) {
            System.out.println("llama_model_load: error loading model: failed to open " + name + ".gguf");
            System.exit(1);
        }
        AtomicReference<String> last = new AtomicReference<>("");
        HttpServer server = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), port), 0);
        server.setExecutor(Executors.newVirtualThreadPerTaskExecutor());
        server.createContext("/health", ex -> send(ex, 200, "{\"status\":\"ok\"}"));
        server.createContext("/last", ex -> send(ex, 200, last.get()));
        server.createContext("/v1/chat/completions", ex -> {
            String body = new String(ex.getRequestBody().readAllBytes(), StandardCharsets.UTF_8);
            last.set(body);
            String prompt = lastContent(body);
            if (prompt.equals("boom")) {
                send(ex, 500, "{\"error\":{\"code\":500,\"message\":\"engine exploded\",\"type\":\"server_error\"}}");
            } else if (body.contains("\"stream\":true")) {
                stream(ex, name, prompt, body.contains("\"tools\""));
            } else if (body.contains("\"tools\"")) {
                send(ex, 200, "{\"choices\":[{\"index\":0,\"finish_reason\":\"tool_calls\",\"message\":{\"role\":\"assistant\",\"content\":\"\","
                        + "\"tool_calls\":[{\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"read_file\",\"arguments\":\"{\\\"path\\\":\\\"pom.xml\\\"}\"}}]}}],"
                        + "\"usage\":{\"prompt_tokens\":20,\"completion_tokens\":8,\"total_tokens\":28}}");
            } else {
                send(ex, 200, "{\"choices\":[{\"index\":0,\"finish_reason\":\"stop\",\"message\":{\"role\":\"assistant\",\"content\":\"fake "
                        + name + ": " + prompt + "\"}}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":4,\"total_tokens\":14}}");
            }
        });
        server.start();
        System.out.println("fake engine " + name + " ready");
    }

    static String lastContent(String body) {
        Matcher m = CONTENT.matcher(body);
        String found = "";
        while (m.find()) found = m.group(1);
        return found;
    }

    static void stream(HttpExchange ex, String name, String prompt, boolean tools) throws IOException {
        ex.getResponseHeaders().set("Content-Type", "text/event-stream");
        ex.sendResponseHeaders(200, 0);
        int count = prompt.equals("long") ? 1000 : 4;
        List<String> pieces = new java.util.ArrayList<>();
        for (String w : ("fake " + name + " says " + prompt).split(" ")) pieces.add(w + " ");
        try (OutputStream out = ex.getResponseBody()) {
            event(out, "{\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":null},\"finish_reason\":null}]}");
            for (int i = 0; i < count; i++) {
                String piece = prompt.equals("long") ? "tok" + i + " " : pieces.get(i % pieces.size());
                event(out, "{\"choices\":[{\"index\":0,\"delta\":{\"content\":\"" + piece + "\"},\"finish_reason\":null}]}");
                Thread.sleep(60);
            }
            event(out, "{\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}");
            event(out, "{\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":" + count + ",\"total_tokens\":" + (10 + count) + "}}");
            event(out, "[DONE]");
        } catch (IOException | InterruptedException e) {
            // the gateway closed the connection
        }
    }

    static void event(OutputStream out, String data) throws IOException {
        out.write(("data: " + data + "\n\n").getBytes(StandardCharsets.UTF_8));
        out.flush();
    }

    static void send(HttpExchange ex, int status, String body) throws IOException {
        byte[] b = body.getBytes(StandardCharsets.UTF_8);
        ex.getResponseHeaders().set("Content-Type", "application/json");
        ex.sendResponseHeaders(status, b.length);
        try (OutputStream out = ex.getResponseBody()) { out.write(b); }
    }
}

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
class GatewayTest {
    static final String KEY = "sk-test";

    @DynamicPropertySource
    static void engines(DynamicPropertyRegistry r) {
        String java = ProcessHandle.current().info().command().orElse("java");
        String cp = System.getProperty("java.class.path");
        r.add("loop.gateway.api-keys[0]", () -> KEY);
        r.add("loop.gateway.max-loaded", () -> 1);
        r.add("loop.gateway.start-timeout", () -> "30s");
        r.add("loop.gateway.tick", () -> "1h");
        String[] names = {"alpha", "beta", "broken"};
        for (int i = 0; i < names.length; i++) {
            String n = names[i], mode = n.equals("broken") ? "broken" : "ok";
            r.add("loop.gateway.models[" + i + "].name", () -> n);
            r.add("loop.gateway.models[" + i + "].command", () -> String.join(",", java, "-cp", cp,
                    "dev.loop.gateway.FakeEngine", "{port}", n, mode));
        }
    }

    @LocalServerPort int port;
    @Autowired BackendManager manager;
    final HttpClient http = HttpClient.newHttpClient();

    @BeforeEach
    void unloadEverything() { manager.unloadIdle(Instant.now().plus(Duration.ofDays(1))); }

    HttpResponse<String> get(String path, String key) throws Exception {
        HttpRequest.Builder b = HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port + path));
        if (key != null) b.header("Authorization", "Bearer " + key);
        return http.send(b.build(), HttpResponse.BodyHandlers.ofString());
    }

    HttpRequest.Builder post(String json) {
        return HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port + "/v1/chat/completions"))
                .header("Authorization", "Bearer " + KEY).header("Content-Type", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(json));
    }

    HttpResponse<String> chat(String json) throws Exception {
        return http.send(post(json).build(), HttpResponse.BodyHandlers.ofString());
    }

    static String user(String model, String text, String extra) {
        return "{\"model\":\"" + model + "\",\"messages\":[{\"role\":\"user\",\"content\":\"" + text + "\"}]" + extra + "}";
    }

    @Test
    void modelsComeFromTheConfiguration() throws Exception {
        assertThat(get("/v1/models", KEY).body()).contains("\"id\":\"alpha\"", "\"id\":\"beta\"");
        assertThat(get("/api/tags", KEY).body()).contains("\"name\":\"alpha\"");
    }
    @Test
    void keysAreChecked() throws Exception {
        assertThat(get("/v1/models", null).statusCode()).isEqualTo(401);
        assertThat(get("/v1/models", "wrong").body()).contains("invalid_api_key");
        assertThat(get("/actuator/health", null).statusCode()).isEqualTo(200);
        assertThat(get("/actuator/metrics", null).statusCode()).isEqualTo(401);
    }
    @Test
    void firstRequestLoadsTheEngine() throws Exception {
        assertThat(get("/api/ps", KEY).body()).isEqualTo("{\"models\":[]}");
        HttpResponse<String> res = chat(user("alpha", "hello", ""));
        assertThat(res.statusCode()).isEqualTo(200);
        assertThat(res.body()).contains("fake alpha: hello");
        assertThat(get("/api/ps", KEY).body()).contains("\"name\":\"alpha\"");
    }
    @Test
    void unknownFieldsAndToolsPassThrough() throws Exception {
        String res = chat(user("alpha", "hi", ",\"top_k\":7,\"custom_field\":{\"a\":1},\"tools\":[{\"type\":\"function\"}]")).body();
        assertThat(res).contains("\"tool_calls\"", "read_file");
        try (BackendManager.Lease lease = manager.acquire("alpha")) {
            String seen = http.send(HttpRequest.newBuilder(URI.create(lease.baseUrl().replace("/v1", "/last"))).build(),
                    HttpResponse.BodyHandlers.ofString()).body();
            assertThat(seen).contains("\"top_k\":7", "\"custom_field\":{\"a\":1}");
        }
    }

    @Test
    void streamingArrivesWhileTheEngineIsStillWriting() throws Exception {
        HttpResponse<java.io.InputStream> res = http.send(post(user("alpha", "stream me", ",\"stream\":true")).build(),
                HttpResponse.BodyHandlers.ofInputStream());
        assertThat(res.headers().firstValue("Content-Type")).hasValue("text/event-stream");
        List<Long> arrivals = new ArrayList<>();
        List<String> data = new ArrayList<>();
        long start = System.nanoTime();
        try (BufferedReader in = new BufferedReader(new InputStreamReader(res.body(), StandardCharsets.UTF_8))) {
            for (String line; (line = in.readLine()) != null; ) {
                if (line.startsWith("data:")) { data.add(line); arrivals.add((System.nanoTime() - start) / 1_000_000); }
            }
        }
        assertThat(data.getLast()).isEqualTo("data: [DONE]");
        assertThat(data).anyMatch(l -> l.contains("\"usage\""));
        long spread = arrivals.getLast() - arrivals.get(1);
        assertThat(spread).as("events are spread over time, not delivered as one block").isGreaterThanOrEqualTo(150);
    }

    @Test
    void aClientThatDisconnectsReleasesTheEngine() throws Exception {
        HttpResponse<java.io.InputStream> res = http.send(post(user("alpha", "long", ",\"stream\":true")).build(),
                HttpResponse.BodyHandlers.ofInputStream());
        BufferedReader in = new BufferedReader(new InputStreamReader(res.body(), StandardCharsets.UTF_8));
        for (int seen = 0; seen < 3; ) if (in.readLine().startsWith("data:")) seen++;
        assertThat(manager.loaded().getFirst().active()).isEqualTo(1);
        in.close();
        long deadline = System.nanoTime() + Duration.ofSeconds(5).toNanos();
        while (manager.loaded().getFirst().active() != 0 && System.nanoTime() < deadline) Thread.sleep(50);
        assertThat(manager.loaded().getFirst().active()).isZero();
    }

    @Test
    void anotherModelEvictsTheIdleOne() throws Exception {
        chat(user("alpha", "one", ""));
        chat(user("beta", "two", ""));
        String ps = get("/api/ps", KEY).body();
        assertThat(ps).contains("\"name\":\"beta\"").doesNotContain("alpha");
    }
    @Test
    void idleModelsAreUnloaded() throws Exception {
        chat(user("alpha", "one", ""));
        manager.unloadIdle(Instant.now().plus(Duration.ofMinutes(1)));
        assertThat(manager.loaded()).hasSize(1);
        manager.unloadIdle(Instant.now().plus(Duration.ofMinutes(6)));
        assertThat(manager.loaded()).isEmpty();
    }

    @Test
    void errorsHaveTheOpenAiShape() throws Exception {
        assertThat(chat(user("nope", "x", "")).statusCode()).isEqualTo(404);
        assertThat(chat(user("nope", "x", "")).body()).contains("model_not_found");
        assertThat(chat("{\"messages\":[]}").body()).contains("The field 'model' is required");
        assertThat(chat("{not json").statusCode()).isEqualTo(400);
        HttpResponse<String> boom = chat(user("alpha", "boom", ""));
        assertThat(boom.statusCode()).isEqualTo(500);
        assertThat(boom.body()).contains("engine exploded");
        HttpResponse<String> broken = chat(user("broken", "x", ""));
        assertThat(broken.statusCode()).isEqualTo(503);
        assertThat(broken.body()).contains("exited with code 1", "failed to open broken.gguf");
    }

public class GatewayChat {
    // ChatClient, Message, Request, Result, Usage, Listener, Cancel, ApiException:
    // the library of Project 11, unchanged.

    public static void main(String[] args) throws Exception {
        String baseUrl = args[0], key = args[1], model = args[2];
        try (ChatClient bad = new ChatClient(baseUrl, "wrong-key")) {
            bad.chat(new Request(model, List.of(Message.user("hi")), null, null, null));
        } catch (ApiException e) {
            System.out.println("wrong key: " + e.status + " " + e.type);
        }
        try (ChatClient client = new ChatClient(baseUrl, key)) {
            List<Message> history = new ArrayList<>(List.of(Message.system("You are a coding assistant."), Message.user("Say hello.")));
            StringBuilder shown = new StringBuilder();
            Result result = client.stream(new Request(model, history, 0.0, args.length > 3 ? Integer.parseInt(args[3]) : null, null),
                    piece -> shown.append('|').append(piece), new Cancel());
            System.out.println("deltas: " + shown + "|");
            System.out.println("finish=" + result.finishReason() + " " + result.usage());
        }
    }
}

public class GatewayAgent {
    // sealed interface Message and its records, ToolSpec, ModelResponse, ModelClient, OpenAiModel:
    // the adapter of the coding-agent lesson (Example 2), unchanged.

    public static void main(String[] args) throws Exception {
        ModelClient client = new OpenAiModel(args[0], args[1]);
        List<ToolSpec> tools = List.of(new ToolSpec("read_file", "Read a file in the workspace.",
                JSON.readTree("{\"type\":\"object\",\"properties\":{\"path\":{\"type\":\"string\"}},\"required\":[\"path\"]}")));
        List<Message> history = new ArrayList<>(List.of(new SystemMessage("You are loop, a coding agent."), new UserMessage("Read pom.xml.")));
        ModelResponse r = client.chat(history, tools);
        if (r.message().toolCalls().isEmpty()) {
            System.out.println("text answer: [" + r.message().text() + "] tokens=" + r.promptTokens() + "+" + r.completionTokens());
            return;
        }
        ToolCall call = r.message().toolCalls().get(0);
        System.out.println("tool call: " + call.name() + " id=" + call.id() + " args=" + JSON.readTree(call.arguments())
                + " tokens=" + r.promptTokens() + "+" + r.completionTokens());
    }
}