feat: extend-TTL, image pre-pull, and session UX

apps/api/src/__tests__/_fakes.ts

@@ -52,6 +52,16 @@ export const createFakeSessionRuntime = (
     async listSessions(): Promise<AirlockSession[]> {
       return current ? [current] : [];
     },
+    async extendSession(sessionId: string, ttlSeconds: number): Promise<AirlockSession | null> {
+      if (!current || current.sessionId !== sessionId) {
+        return null;
+      }
+      current = { ...current, expiresAt: new Date(Date.now() + ttlSeconds * 1000).toISOString() };
+      return current;
+    },
+    async pullBrowserImages() {
+      return [{ image: "kasmweb/chromium:1.18.0", ok: true }];
+    },
     async stopSession(sessionId: string): Promise<boolean> {
       stopped.push(sessionId);
       const wasPresent = current !== null;

apps/api/src/__tests__/app.test.ts

@@ -108,6 +108,48 @@ describe("Airlock API", () => {
     expect(response.status).toBe(429);
   });
 
+  it("extends a session's TTL via PATCH", async () => {
+    const app = createApp({
+      config: testConfig,
+      sessionRuntime: createFakeSessionRuntime()
+    });
+
+    const response = await request(app).patch("/api/sessions/session-1").send({ ttlSeconds: 600 });
+    expect(response.status).toBe(200);
+    expect(response.body.sessionId).toBe("session-1");
+  });
+
+  it("returns 404 when extending an unknown session", async () => {
+    const app = createApp({
+      config: testConfig,
+      sessionRuntime: createFakeSessionRuntime({ initial: null })
+    });
+
+    const response = await request(app).patch("/api/sessions/nope").send({ ttlSeconds: 600 });
+    expect(response.status).toBe(404);
+  });
+
+  it("rejects an invalid extend body", async () => {
+    const app = createApp({
+      config: testConfig,
+      sessionRuntime: createFakeSessionRuntime()
+    });
+
+    const response = await request(app).patch("/api/sessions/session-1").send({ ttlSeconds: 5 });
+    expect(response.status).toBe(400);
+  });
+
+  it("pulls browser images on demand", async () => {
+    const app = createApp({
+      config: testConfig,
+      sessionRuntime: createFakeSessionRuntime()
+    });
+
+    const response = await request(app).post("/api/images/pull");
+    expect(response.status).toBe(200);
+    expect(Array.isArray(response.body.images)).toBe(true);
+  });
+
   it("serves Prometheus metrics", async () => {
     const app = createApp({
       config: testConfig,

apps/api/src/__tests__/docker-session-runtime.test.ts

@@ -62,6 +62,7 @@ class FakeDocker {
   createCalls: CreateContainerCall[] = [];
   startedIds: string[] = [];
   networksCreated: string[] = [];
+  pulledImages: string[] = [];
   pingCount = 0;
   private readonly hostPortByContainerPort: Record<string, string>;
   private readonly nextContainerId: string;
@@ -139,6 +140,21 @@ class FakeDocker {
     this.pingCount += 1;
     return "OK";
   };
+
+  getImage = (_name: string) => ({
+    inspect: async (): Promise<unknown> => ({ Id: "img" })
+  });
+
+  pull = async (image: string): Promise<unknown> => {
+    this.pulledImages.push(image);
+    return {};
+  };
+
+  modem = {
+    followProgress: (_stream: unknown, onFinished: (error: Error | null) => void): void => {
+      onFinished(null);
+    }
+  };
 }
 
 describe("DockerSessionRuntime", () => {
@@ -358,6 +374,48 @@ describe("DockerSessionRuntime", () => {
     expect(docker.pingCount).toBe(1);
   });
 
+  it("extendSession pushes the expiry out and keeps the container from pruning", async () => {
+    const docker = new FakeDocker({ containers: [makeFakeContainer()] });
+    const runtime = new DockerSessionRuntime({
+      config: baseConfig,
+      docker: docker as unknown as Docker
+    });
+
+    // The label expiry is 2026-04-30T12:30:00Z; extend well past a later "now".
+    const extended = await runtime.extendSession("s-1", 3600);
+    expect(extended).not.toBeNull();
+    expect(new Date(extended!.expiresAt).getTime()).toBeGreaterThan(
+      new Date("2026-04-30T13:00:00.000Z").getTime()
+    );
+
+    // At a time past the original label expiry, the extended session survives.
+    const pruned = await runtime.pruneExpiredSessions(new Date("2026-04-30T12:45:00.000Z"));
+    expect(pruned).toBe(0);
+    expect(docker.removed).toEqual([]);
+  });
+
+  it("extendSession returns null for an unknown session", async () => {
+    const docker = new FakeDocker({ containers: [makeFakeContainer()] });
+    const runtime = new DockerSessionRuntime({
+      config: baseConfig,
+      docker: docker as unknown as Docker
+    });
+    expect(await runtime.extendSession("missing", 600)).toBeNull();
+  });
+
+  it("pullBrowserImages pulls each unique configured image", async () => {
+    const docker = new FakeDocker();
+    const runtime = new DockerSessionRuntime({
+      config: baseConfig,
+      docker: docker as unknown as Docker
+    });
+
+    const results = await runtime.pullBrowserImages();
+    expect(results.every((r) => r.ok)).toBe(true);
+    expect(docker.pulledImages).toContain("kasmweb/chromium:1.18.0");
+    expect(docker.pulledImages).toContain("kasmweb/tor-browser:1.18.0");
+  });
+
   it("createSession removes the container and throws when no host port is mapped", async () => {
     const docker = new FakeDocker({
       nextContainerId: "doomed",

apps/api/src/app.ts

@@ -16,7 +16,7 @@ import { logger } from "./logger";
 import { Metrics } from "./metrics";
 import { createRateLimit } from "./rate-limit";
 import { resolveOrRespond } from "./resolve-session";
-import { createSessionBodySchema } from "./schemas";
+import { createSessionBodySchema, extendSessionBodySchema } from "./schemas";
 import { toSessionResponse } from "./session-response";
 
 export interface CreateAppOptions {
@@ -183,6 +183,43 @@ export const createApp = ({ config, sessionRuntime, metrics }: CreateAppOptions)
     })
   );
 
+  app.patch(
+    "/api/sessions/:sessionId",
+    bearerAuth,
+    asyncRoute(async (request: Request, response: Response) => {
+      const parsed = extendSessionBodySchema.safeParse(request.body);
+      if (!parsed.success) {
+        response.status(400).json({
+          error: parsed.error.issues.map((issue) => issue.message).join(", ")
+        });
+        return;
+      }
+
+      const session = await sessionRuntime.extendSession(
+        request.params.sessionId,
+        parsed.data.ttlSeconds
+      );
+      if (!session) {
+        response.status(404).json({ error: "Session not found." });
+        return;
+      }
+      logger.info("session.extended", {
+        sessionId: session.sessionId,
+        expiresAt: session.expiresAt
+      });
+      response.json(toSessionResponse(session, config));
+    })
+  );
+
+  app.post(
+    "/api/images/pull",
+    bearerAuth,
+    asyncRoute(async (_request: Request, response: Response) => {
+      const images = await sessionRuntime.pullBrowserImages();
+      response.json({ images });
+    })
+  );
+
   app.delete(
     "/api/sessions/:sessionId",
     bearerAuth,

apps/api/src/docker-session-runtime.ts

@@ -6,8 +6,10 @@ import {
   AirlockSession,
   CreateSessionInput,
   DecodedSessionLabels,
+  ImagePullResult,
   SessionRuntime,
   browserProfile,
+  clampTtl,
   decodeSessionLabels,
   encodeSessionLabels,
   expiresAt as computeExpiresAt,
@@ -86,6 +88,11 @@ export interface DockerSessionRuntimeOptions {
 export class DockerSessionRuntime implements SessionRuntime {
   private readonly docker: Docker;
   private readonly config: AirlockConfig;
+  // Extended expiries by sessionId. Docker labels are immutable after creation
+  // and there is no database, so a TTL extension is held in memory and applied
+  // on read/prune. It is best-effort: an API restart drops it and the session
+  // reverts to its label expiry. Single-node only.
+  private readonly expiryOverrides = new Map<string, string>();
 
   constructor(options: DockerSessionRuntimeOptions) {
     this.config = options.config;
@@ -110,6 +117,9 @@ export class DockerSessionRuntime implements SessionRuntime {
     if (launch.networkIsolation) {
       await this.ensureNetwork(launch.networkName);
     }
+    // createContainer does not auto-pull; make sure the image is present so the
+    // first launch on a fresh host works instead of failing with "no such image".
+    await this.ensureImage(image);
 
     const launchEnv = {
       ...profile.buildLaunchEnv({
@@ -183,6 +193,59 @@ export class DockerSessionRuntime implements SessionRuntime {
     }
   }
 
+  async extendSession(sessionId: string, ttlSeconds: number): Promise<AirlockSession | null> {
+    const match = await this.findContainerBySessionId(sessionId);
+    if (!match) {
+      return null;
+    }
+    const newExpiry = computeExpiresAt(new Date(), clampTtl(ttlSeconds));
+    this.expiryOverrides.set(sessionId, newExpiry.toISOString());
+    return this.mapContainerToSession(match.container, match.decoded);
+  }
+
+  async pullBrowserImages(): Promise<ImagePullResult[]> {
+    const images = [...new Set(Object.values(this.config.containerLaunch.browserImages))];
+    const results: ImagePullResult[] = [];
+    for (const image of images) {
+      try {
+        await this.pullImage(image);
+        results.push({ image, ok: true });
+      } catch (error) {
+        logger.warn("image.pull_failed", {
+          image,
+          message: error instanceof Error ? error.message : String(error)
+        });
+        results.push({ image, ok: false });
+      }
+    }
+    return results;
+  }
+
+  private async ensureImage(image: string): Promise<void> {
+    try {
+      await this.docker.getImage(image).inspect();
+      return;
+    } catch (error) {
+      if (!isContainerNotFound(error)) {
+        throw error;
+      }
+    }
+    await this.pullImage(image);
+  }
+
+  private async pullImage(image: string): Promise<void> {
+    const stream = await this.docker.pull(image);
+    await new Promise<void>((resolve, reject) => {
+      this.docker.modem.followProgress(stream, (error: Error | null) => {
+        if (error) {
+          reject(error);
+        } else {
+          resolve();
+        }
+      });
+    });
+  }
+
   async getSession(sessionId: string): Promise<AirlockSession | null> {
     const match = await this.findContainerBySessionId(sessionId);
     if (!match) {
@@ -215,24 +278,37 @@ export class DockerSessionRuntime implements SessionRuntime {
       return false;
     }
     await this.safeRemoveContainer(match.container.Id);
+    this.expiryOverrides.delete(sessionId);
     return true;
   }
 
   async pruneExpiredSessions(now: Date = new Date()): Promise<number> {
     const containers = await this.listManagedContainers(true);
+    const liveSessionIds = new Set<string>();
     let pruned = 0;
 
     for (const container of containers) {
       const decoded = decodeSessionLabels(container.Labels, this.config.sessionDefaults.browser);
       if (!decoded) {
         continue;
       }
+      liveSessionIds.add(decoded.sessionId);
 
-      if (isExpired(decoded, now)) {
+      // Honor an in-memory TTL extension over the label's original expiry.
+      const effectiveExpiresAt = this.expiryOverrides.get(decoded.sessionId) ?? decoded.expiresAt;
+      if (isExpired({ expiresAt: effectiveExpiresAt }, now)) {
         await this.safeRemoveContainer(container.Id);
+        this.expiryOverrides.delete(decoded.sessionId);
         pruned += 1;
       }
     }
+
+    // Drop overrides for sessions that no longer exist (e.g. AutoRemoved).
+    for (const sessionId of this.expiryOverrides.keys()) {
+      if (!liveSessionIds.has(sessionId)) {
+        this.expiryOverrides.delete(sessionId);
+      }
+    }
     return pruned;
   }
 
@@ -275,7 +351,8 @@ export class DockerSessionRuntime implements SessionRuntime {
       browserUrl: profile.streamUrl(this.config.server.sessionHost, hostPort),
       vncPassword: decoded.vncPassword,
       createdAt: decoded.createdAt,
-      expiresAt: decoded.expiresAt
+      // Reflect an in-memory TTL extension when one is set for this session.
+      expiresAt: this.expiryOverrides.get(decoded.sessionId) ?? decoded.expiresAt
     };
   }
 

apps/api/src/schemas.ts

@@ -15,3 +15,9 @@ export const createSessionBodySchema = z.object({
 });
 
 export type CreateSessionBody = z.infer<typeof createSessionBodySchema>;
+
+export const extendSessionBodySchema = z.object({
+  ttlSeconds: z.number().int().min(TTL_MIN_SECONDS).max(TTL_MAX_SECONDS)
+});
+
+export type ExtendSessionBody = z.infer<typeof extendSessionBodySchema>;