fix(server): keep Cue skipped across serving restarts (#271)

apps/server/src/contextmd/ContextmdRoute.test.ts

@@ -233,13 +233,21 @@ const servedSession = (() => {
 /** Serve the contextmd routes on a test server over a migrated database prepared by `setup`. */
 const serve = <R>(
   setup: Effect.Effect<void, never, R | SqlClient.SqlClient>,
-  options: { readonly breakStateFile?: boolean; readonly cueUrl?: string } = {},
+  options: {
+    readonly breakStateFile?: boolean;
+    readonly cueUrl?: string;
+    /** A serving database already marked, so this server starts as the next process would. */
+    readonly seedState?: string;
+  } = {},
 ) =>
   Effect.gen(function* () {
     const configContext = yield* Layer.build(
       ServerConfig.layerTest(process.cwd(), { prefix: "contextmd-route-test-" }),
     );
     const config = Context.get(configContext, ServerConfig.ServerConfig);
+    if (options.seedState !== undefined) {
+      NodeFS.copyFileSync(options.seedState, NodePath.join(config.stateDir, "contextmd.sqlite"));
+    }
     const store = writeStore(config.baseDir);
     if (options.cueUrl !== undefined) {
       const { cue } = writeCue(config.baseDir);
@@ -410,6 +418,7 @@ const serve = <R>(
           Effect.orElseSucceed(() => "reset" as const),
         );
     };
+    const sqlitePath = NodePath.join(config.stateDir, "contextmd.sqlite");
     return {
       moment,
       settle,
@@ -422,6 +431,7 @@ const serve = <R>(
       outcomes,
       database,
       cueLog,
+      sqlitePath,
     };
   });
 
@@ -994,6 +1004,49 @@ describe("contextmd routes", () => {
     }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
   );
 
+  it.live("a restarted server keeps Cue skipped until its health document is back", () =>
+    Effect.gen(function* () {
+      const url = yield* Effect.promise(closedPort);
+      const setup = Effect.all([
+        insertAttempt({ attempt: "a", run: "run-1", status: "running", native: servedSession }),
+        projectUserMessage("Ship it"),
+      ]).pipe(Effect.asVoid);
+      const first = yield* serve(setup, { cueUrl: url });
+      assert.strictEqual(
+        (yield* first.moment({
+          hook_event_name: "UserPromptSubmit",
+          session_id: servedSession,
+          cwd: "/tmp",
+          prompt: "Ship it",
+        })).status,
+        204,
+      );
+      assert.strictEqual(first.cueLog().moments[0]!.cue_status, "connect_failed");
+      const db = new NodeSqlite.DatabaseSync(first.sqlitePath);
+      db.exec("PRAGMA wal_checkpoint(FULL)");
+      db.close();
+      const seed = `${first.sqlitePath}.seed`;
+      NodeFS.copyFileSync(first.sqlitePath, seed);
+      const second = yield* serve(setup, { cueUrl: url, seedState: seed });
+      yield* second.moment(
+        {
+          hook_event_name: "UserPromptSubmit",
+          session_id: servedSession,
+          cwd: "/tmp",
+          prompt: "Ship it",
+        },
+        { delivery: "second-user" },
+      );
+      yield* second.moment(claude("PreToolUse", servedSession, "echo a"), {
+        delivery: "second-pre",
+      });
+      assert.deepStrictEqual(
+        second.cueLog().moments.map((m) => m.cue_status),
+        ["connect_failed", "not_called", "not_called"],
+      );
+    }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
+  );
+
   it.live("keeps the current user message in Cue's window after hundreds of tool calls", () =>
     Effect.gen(function* () {
       const cue = yield* fakeCue({ "agentsmd-research": 2.5 });

apps/server/src/contextmd/ContextmdServing.ts

@@ -16,7 +16,9 @@
  * With Cue configured (`cue` entry directories and `cue_url`), a joined moment also asks the Cue service
  * (`Cue.ts`): a user message and a failed call score at once, a pre-tool moment starts a score its tool
  * result collects. The window items come from the thread's projected turn items. Cue never holds the
- * exact-key lane up: an unreachable service is skipped machine-wide until a probe sees it again.
+ * exact-key lane up: an unreachable service is skipped machine-wide until a probe sees its health
+ * document again. The skip is stored in the serving database, so a restarted server does not call
+ * Cue on every new process while the service is down.
  *
  * @module ContextmdServing
  */
@@ -208,14 +210,20 @@ const make = Effect.gen(function* () {
     });
   }
   const { db, state } = opened.value;
+  const clock = yield* Clock.Clock;
   const parsed: ParseCache = new Map();
   const cueParsed: ParseCache = new Map();
   // One client per Cue service URL, so its availability is shared by every moment on this machine.
   const cueClients = new Map<string, CueClient>();
   const cueClient = (url: string) => {
     let client = cueClients.get(url);
     if (client === undefined) {
-      client = new CueClient(httpTransport(url));
+      // A previous process may already have seen this URL fail. Start skipped until health returns.
+      const nowMs = () => clock.currentTimeMillisUnsafe();
+      client = new CueClient(httpTransport(url), nowMs, !state.cueDown(url), (available) => {
+        if (available) state.clearCueDown(url);
+        else state.markCueDown(url, nowMs() / 1000);
+      });
       cueClients.set(url, client);
     }
     return client;

apps/server/src/contextmd/Cue.test.ts

@@ -58,19 +58,38 @@ describe("Cue client", () => {
     const { transport, calls } = scripted([
       { kind: "connect_failed" },
       { kind: "connect_failed" },
-      { kind: "reply", status: 200, body: "{}" },
+      { kind: "reply", status: 200, body: "ok" },
+      { kind: "reply", status: 200, body: JSON.stringify({ ok: true, adapter: "v1" }) },
       { kind: "reply", status: 200, body: OK },
     ]);
-    const c = new CueClient(transport);
+    const changes: boolean[] = [];
+    const c = new CueClient(transport, Date.now, true, (up) => changes.push(up));
     assert.strictEqual((await c.score(REQ, WAIT)).status, "connect_failed");
     assert.isFalse(c.isAvailable());
+    assert.deepStrictEqual(changes, [false]);
     assert.strictEqual((await c.score(REQ, WAIT)).status, "not_called");
     assert.strictEqual(c.start("k", REQ).status, "not_called");
     assert.strictEqual((await c.collect("k")).status, "not_called");
     assert.strictEqual((await c.collect("no-start-while-down")).status, "not_called");
     assert.strictEqual(calls.length, 1);
     assert.isFalse(await c.probe());
+    assert.isFalse(await c.probe());
+    assert.isTrue(await c.probe());
+    assert.deepStrictEqual(changes, [false, true]);
+    assert.strictEqual((await c.score(REQ, WAIT)).status, "ok");
+  });
+
+  it("stays skipped when constructed down, until a health document arrives", async () => {
+    const { transport, calls } = scripted([
+      { kind: "reply", status: 200, body: JSON.stringify({ ok: true, adapter: "v1" }) },
+      { kind: "reply", status: 200, body: OK },
+    ]);
+    const c = new CueClient(transport, Date.now, false);
+    assert.strictEqual((await c.score(REQ, WAIT)).status, "not_called");
+    assert.strictEqual(c.start("k", REQ).status, "not_called");
+    assert.strictEqual(calls.length, 0);
     assert.isTrue(await c.probe());
+    assert.strictEqual(calls.at(-1)!.path, "/health");
     assert.strictEqual((await c.score(REQ, WAIT)).status, "ok");
   });
 

apps/server/src/contextmd/Cue.ts

@@ -6,8 +6,8 @@
  * - A `pre_tool` moment only starts a request; the call's `tool_result` collects it (latency hidden behind the
  *   tool run). A `user_message` or a failed call scores at once and waits what is left of the moment's budget.
  * - A connect failure marks the service unavailable for this whole environment (the machine); later moments
- *   skip Cue (`not_called`) until a probe every `PROBE_SECONDS` sees `/health` answer. A slow reply is
- *   `timeout` for that call only.
+ *   skip Cue (`not_called`) until a probe every `PROBE_SECONDS` gets the Cue health document (`ok` and
+ *   `adapter`). A slow reply is `timeout` for that call only.
  * - A 4xx reply is `rejected`: the service refused this request (its reason is logged), the service stays
  *   available.
  * - A collect with no request for that call is `lost` (a restart dropped it), or `not_called` while the
@@ -181,16 +181,45 @@ const parseScore = (status: number, body: string, ms: number): CueReply => {
   }
 };
 
+/** What `/health` returns while Cue is actually serving. A status line alone is not enough. */
+const healthDocument = (body: string): boolean => {
+  try {
+    const v = JSON.parse(body) as Record<string, unknown>;
+    return v["ok"] === true && typeof v["adapter"] === "string" && v["adapter"] !== "";
+  } catch {
+    return false;
+  }
+};
+
 export class CueClient {
-  private available = true;
+  private available: boolean;
   private probing = false;
   private readonly pending = new Map<string, Pending>();
   private readonly transport: Transport;
   private readonly clock: () => number;
+  private readonly onAvailability: ((available: boolean) => void) | undefined;
 
-  constructor(transport: Transport, clock: () => number = Date.now) {
+  constructor(
+    transport: Transport,
+    clock: () => number = Date.now,
+    initiallyAvailable = true,
+    onAvailability?: (available: boolean) => void,
+  ) {
     this.transport = transport;
     this.clock = clock;
+    this.available = initiallyAvailable;
+    this.onAvailability = onAvailability;
+  }
+
+  private setAvailable(available: boolean) {
+    if (this.available === available) return;
+    this.available = available;
+    // Recording the belief is best-effort: a disk error must not fail the moment.
+    try {
+      this.onAvailability?.(available);
+    } catch {
+      return;
+    }
   }
 
   /** Whether the service is believed reachable (false between a connect failure and a good probe). */
@@ -213,25 +242,21 @@ export class CueClient {
       .then((r) => {
         const ms = this.clock() - t0;
         if (r.kind !== "reply") {
-          if (r.kind === "connect_failed") this.markUnavailable();
+          if (r.kind === "connect_failed") this.setAvailable(false);
           return { status: r.kind, ms, waitMs: null } satisfies CueReply;
         }
         return parseScore(r.status, r.body, ms);
       })
       .catch((): CueReply => ({ status: "error", ms: this.clock() - t0, waitMs: null }));
   }
 
-  private markUnavailable() {
-    this.available = false;
-  }
-
-  /** One health probe; marks the service available again when it answers. */
+  /** One health probe; marks the service available again only for a Cue health document. */
   async probe(): Promise<boolean> {
     if (this.probing) return this.available;
     this.probing = true;
     try {
       const r = await this.transport("GET", "/health", null, 2000);
-      if (r.kind === "reply" && r.status === 200) this.available = true;
+      if (r.kind === "reply" && r.status === 200 && healthDocument(r.body)) this.setAvailable(true);
       return this.available;
     } catch {
       return this.available;

apps/server/src/contextmd/ServingState.test.ts

@@ -1068,6 +1068,27 @@ describe("ServingState Cue", () => {
     );
   });
 
+  it("remembers that Cue is down in the database the next process opens", () => {
+    const dir = NodeFS.mkdtempSync(NodePath.join(NodeOS.tmpdir(), "cue-down-"));
+    const file = NodePath.join(dir, "contextmd.sqlite");
+    try {
+      const first = new ServingState(new NodeSqlite.DatabaseSync(file));
+      assert.isFalse(first.cueDown("http://127.0.0.1:9"));
+      first.markCueDown("http://127.0.0.1:9", 10);
+      first.markCueDown("http://127.0.0.1:9", 11);
+      assert.isTrue(first.cueDown("http://127.0.0.1:9"));
+      assert.isFalse(first.cueDown("http://127.0.0.1:10"));
+      const again = new ServingState(new NodeSqlite.DatabaseSync(file));
+      assert.isTrue(again.cueDown("http://127.0.0.1:9"));
+      again.clearCueDown("http://127.0.0.1:9");
+      assert.isFalse(
+        new ServingState(new NodeSqlite.DatabaseSync(file)).cueDown("http://127.0.0.1:9"),
+      );
+    } finally {
+      NodeFS.rmSync(dir, { recursive: true, force: true });
+    }
+  });
+
   it("records an unreachable or slow Cue per moment and still serves exact keys", () => {
     const h = makeHarness();
     const failed = (status: CueReply["status"]): CueMoment => ({

apps/server/src/contextmd/ServingState.ts

@@ -77,6 +77,7 @@ CREATE TABLE IF NOT EXISTS candidates (id INTEGER PRIMARY KEY, moment_id INTEGER
 CREATE INDEX IF NOT EXISTS candidates_moment ON candidates(moment_id);
 CREATE INDEX IF NOT EXISTS candidates_entry ON candidates(entry, served_at);
 CREATE TABLE IF NOT EXISTS heartbeats (at REAL PRIMARY KEY);
+CREATE TABLE IF NOT EXISTS cue_down (url TEXT PRIMARY KEY, since REAL NOT NULL);
 `;
 
 /** Columns added after a table first shipped: created on open when an older state file lacks them. */
@@ -324,6 +325,20 @@ export class ServingState {
     this.run("INSERT OR IGNORE INTO heartbeats (at) VALUES (?)", now);
   }
 
+  /** Whether this machine last saw Cue's connect fail and has not had a health document since. */
+  cueDown(url: string): boolean {
+    return this.get<{ url: string }>("SELECT url FROM cue_down WHERE url = ?", url) !== undefined;
+  }
+
+  /** Remember a connect failure so the next process skips Cue until health returns. */
+  markCueDown(url: string, now: number) {
+    this.run("INSERT OR REPLACE INTO cue_down (url, since) VALUES (?, ?)", url, now);
+  }
+
+  clearCueDown(url: string) {
+    this.run("DELETE FROM cue_down WHERE url = ?", url);
+  }
+
   /** Log a moment that was not served for a reason outside selection. */
   logUnserved(u: Unserved) {
     this.tx(() => {