Add a host-managed wait tool for Codex sessions

apps/server/src/provider/Layers/CodexDynamicTools.test.ts

@@ -0,0 +1,207 @@
+import * as NodeAssert from "node:assert/strict";
+
+import { it } from "@effect/vitest";
+import * as Deferred from "effect/Deferred";
+import * as Effect from "effect/Effect";
+import * as Fiber from "effect/Fiber";
+import * as TestClock from "effect/testing/TestClock";
+import { describe } from "vite-plus/test";
+import type * as EffectCodexSchema from "effect-codex-app-server/schema";
+
+import {
+  handleT3CodexDynamicToolCall,
+  makeT3CodexDynamicToolWaitRegistry,
+  T3_CODEX_DYNAMIC_TOOLS,
+  T3_CODEX_DYNAMIC_TOOL_NAMESPACE,
+  T3_CODEX_WAIT_MAX_DURATION_MS,
+  T3_CODEX_WAIT_MIN_DURATION_MS,
+  T3_CODEX_WAIT_TOOL_NAME,
+} from "./CodexDynamicTools.ts";
+
+function waitCall(
+  arguments_: unknown,
+  overrides: Partial<EffectCodexSchema.DynamicToolCallParams> = {},
+): EffectCodexSchema.DynamicToolCallParams {
+  return {
+    arguments: arguments_,
+    callId: "wait-call-1",
+    namespace: T3_CODEX_DYNAMIC_TOOL_NAMESPACE,
+    threadId: "provider-thread-1",
+    tool: T3_CODEX_WAIT_TOOL_NAME,
+    turnId: "provider-turn-1",
+    ...overrides,
+  };
+}
+
+describe("T3 Codex dynamic tools", () => {
+  it("publishes a bounded namespaced wait tool", () => {
+    const namespace = T3_CODEX_DYNAMIC_TOOLS[0];
+    const wait = namespace?.tools[0];
+
+    NodeAssert.equal(namespace?.type, "namespace");
+    NodeAssert.equal(namespace?.name, T3_CODEX_DYNAMIC_TOOL_NAMESPACE);
+    NodeAssert.equal(wait?.type, "function");
+    NodeAssert.equal(wait?.name, T3_CODEX_WAIT_TOOL_NAME);
+    NodeAssert.deepStrictEqual(wait?.inputSchema, {
+      type: "object",
+      properties: {
+        durationMs: {
+          type: "integer",
+          minimum: T3_CODEX_WAIT_MIN_DURATION_MS,
+          maximum: T3_CODEX_WAIT_MAX_DURATION_MS,
+          description: "How long T3 should wait, in milliseconds.",
+        },
+        reason: {
+          type: "string",
+          maxLength: 500,
+          description: "Optional short explanation of the external work being awaited.",
+        },
+      },
+      required: ["durationMs"],
+      additionalProperties: false,
+    });
+  });
+
+  it.effect("does not complete before the requested duration", () =>
+    Effect.gen(function* () {
+      const fiber = yield* handleT3CodexDynamicToolCall(
+        waitCall({ durationMs: T3_CODEX_WAIT_MIN_DURATION_MS, reason: "CI is running" }),
+      ).pipe(Effect.forkChild);
+
+      yield* TestClock.adjust(T3_CODEX_WAIT_MIN_DURATION_MS - 1);
+      NodeAssert.equal(fiber.pollUnsafe(), undefined);
+
+      yield* TestClock.adjust(1);
+      NodeAssert.deepStrictEqual(yield* Fiber.join(fiber), {
+        success: true,
+        contentItems: [
+          {
+            type: "inputText",
+            text: `Wait completed after ${T3_CODEX_WAIT_MIN_DURATION_MS} ms.`,
+          },
+        ],
+      });
+    }).pipe(Effect.provide(TestClock.layer())),
+  );
+
+  it.effect("cancels an active wait", () =>
+    Effect.gen(function* () {
+      const cancelled = yield* Deferred.make<void>();
+      const fiber = yield* handleT3CodexDynamicToolCall(
+        waitCall({ durationMs: T3_CODEX_WAIT_MAX_DURATION_MS }),
+        Deferred.await(cancelled),
+      ).pipe(Effect.forkChild);
+
+      yield* Effect.yieldNow;
+      yield* Deferred.succeed(cancelled, undefined);
+
+      NodeAssert.deepStrictEqual(yield* Fiber.join(fiber), {
+        success: false,
+        contentItems: [
+          {
+            type: "inputText",
+            text: "Wait cancelled because the turn was interrupted or the session closed.",
+          },
+        ],
+      });
+    }).pipe(Effect.provide(TestClock.layer())),
+  );
+
+  it.effect("cancels waits by turn and cancels all remaining waits on shutdown", () =>
+    Effect.gen(function* () {
+      const registry = yield* makeT3CodexDynamicToolWaitRegistry();
+      const first = yield* registry
+        .handle(
+          waitCall(
+            { durationMs: T3_CODEX_WAIT_MAX_DURATION_MS },
+            { callId: "wait-call-1", turnId: "provider-turn-1" },
+          ),
+        )
+        .pipe(Effect.forkChild);
+      const second = yield* registry
+        .handle(
+          waitCall(
+            { durationMs: T3_CODEX_WAIT_MAX_DURATION_MS },
+            { callId: "wait-call-2", turnId: "provider-turn-2" },
+          ),
+        )
+        .pipe(Effect.forkChild);
+
+      yield* Effect.yieldNow;
+      yield* registry.cancelTurn("provider-turn-1");
+
+      NodeAssert.equal((yield* Fiber.join(first)).success, false);
+      NodeAssert.equal(second.pollUnsafe(), undefined);
+
+      yield* registry.cancelAll;
+      NodeAssert.equal((yield* Fiber.join(second)).success, false);
+    }).pipe(Effect.provide(TestClock.layer())),
+  );
+
+  it.effect("immediately cancels waits registered after their turn was interrupted", () =>
+    Effect.gen(function* () {
+      const registry = yield* makeT3CodexDynamicToolWaitRegistry();
+      yield* registry.cancelTurn("provider-turn-1");
+
+      const response = yield* registry.handle(
+        waitCall({ durationMs: T3_CODEX_WAIT_MAX_DURATION_MS }, { turnId: "provider-turn-1" }),
+      );
+
+      NodeAssert.equal(response.success, false);
+      NodeAssert.match(
+        response.contentItems[0]?.type === "inputText" ? response.contentItems[0].text : "",
+        /Wait cancelled/,
+      );
+    }).pipe(Effect.provide(TestClock.layer())),
+  );
+
+  it.effect("immediately cancels waits registered after session shutdown", () =>
+    Effect.gen(function* () {
+      const registry = yield* makeT3CodexDynamicToolWaitRegistry();
+      yield* registry.cancelAll;
+
+      const response = yield* registry.handle(
+        waitCall({ durationMs: T3_CODEX_WAIT_MAX_DURATION_MS }),
+      );
+
+      NodeAssert.equal(response.success, false);
+      NodeAssert.match(
+        response.contentItems[0]?.type === "inputText" ? response.contentItems[0].text : "",
+        /Wait cancelled/,
+      );
+    }).pipe(Effect.provide(TestClock.layer())),
+  );
+
+  it.effect("rejects malformed and out-of-range arguments without sleeping", () =>
+    Effect.gen(function* () {
+      for (const arguments_ of [
+        {},
+        { durationMs: T3_CODEX_WAIT_MIN_DURATION_MS - 1 },
+        { durationMs: T3_CODEX_WAIT_MAX_DURATION_MS + 1 },
+        { durationMs: 1_000.5 },
+        { durationMs: "1000" },
+      ]) {
+        const response = yield* handleT3CodexDynamicToolCall(waitCall(arguments_));
+        NodeAssert.equal(response.success, false);
+        NodeAssert.match(
+          response.contentItems[0]?.type === "inputText" ? response.contentItems[0].text : "",
+          /Invalid wait arguments/,
+        );
+      }
+    }),
+  );
+
+  it.effect("rejects calls outside the T3 namespace", () =>
+    Effect.gen(function* () {
+      const response = yield* handleT3CodexDynamicToolCall(
+        waitCall({ durationMs: T3_CODEX_WAIT_MIN_DURATION_MS }, { namespace: "other" }),
+      );
+
+      NodeAssert.equal(response.success, false);
+      NodeAssert.match(
+        response.contentItems[0]?.type === "inputText" ? response.contentItems[0].text : "",
+        /Unknown T3 Code dynamic tool/,
+      );
+    }),
+  );
+});

apps/server/src/provider/Layers/CodexDynamicTools.ts

@@ -0,0 +1,173 @@
+import * as Deferred from "effect/Deferred";
+import * as Duration from "effect/Duration";
+import * as Effect from "effect/Effect";
+import * as Ref from "effect/Ref";
+import * as Result from "effect/Result";
+import * as Schema from "effect/Schema";
+import type * as EffectCodexSchema from "effect-codex-app-server/schema";
+
+export const T3_CODEX_DYNAMIC_TOOL_NAMESPACE = "t3";
+export const T3_CODEX_WAIT_TOOL_NAME = "wait";
+export const T3_CODEX_WAIT_MIN_DURATION_MS = 1_000;
+export const T3_CODEX_WAIT_MAX_DURATION_MS = 3_600_000;
+
+const T3CodexWaitArguments = Schema.Struct({
+  durationMs: Schema.Int.check(Schema.isGreaterThanOrEqualTo(T3_CODEX_WAIT_MIN_DURATION_MS)).check(
+    Schema.isLessThanOrEqualTo(T3_CODEX_WAIT_MAX_DURATION_MS),
+  ),
+  reason: Schema.optionalKey(Schema.String.check(Schema.isMaxLength(500))),
+});
+const decodeT3CodexWaitArguments = Schema.decodeUnknownEffect(T3CodexWaitArguments);
+
+export const T3_CODEX_DYNAMIC_TOOLS = [
+  {
+    type: "namespace",
+    name: T3_CODEX_DYNAMIC_TOOL_NAMESPACE,
+    description: "T3 Code host lifecycle tools that can pause work without model polling.",
+    tools: [
+      {
+        type: "function",
+        name: T3_CODEX_WAIT_TOOL_NAME,
+        description:
+          "Wait inside T3 Code without using model inference. Use this for an intentional idle period while external work is running, then check the work once after the wait completes. The wait only lasts while this live T3 provider session remains connected; it is cancelled if the turn is interrupted or the session closes.",
+        inputSchema: {
+          type: "object",
+          properties: {
+            durationMs: {
+              type: "integer",
+              minimum: T3_CODEX_WAIT_MIN_DURATION_MS,
+              maximum: T3_CODEX_WAIT_MAX_DURATION_MS,
+              description: "How long T3 should wait, in milliseconds.",
+            },
+            reason: {
+              type: "string",
+              maxLength: 500,
+              description: "Optional short explanation of the external work being awaited.",
+            },
+          },
+          required: ["durationMs"],
+          additionalProperties: false,
+        },
+      },
+    ],
+  },
+] as const satisfies ReadonlyArray<EffectCodexSchema.V2ThreadStartParams__DynamicToolSpec>;
+
+function textResponse(success: boolean, text: string): EffectCodexSchema.DynamicToolCallResponse {
+  return {
+    success,
+    contentItems: [{ type: "inputText", text }],
+  };
+}
+
+export const handleT3CodexDynamicToolCall = Effect.fn("handleT3CodexDynamicToolCall")(function* (
+  payload: EffectCodexSchema.DynamicToolCallParams,
+  cancelled: Effect.Effect<void> = Effect.never,
+): Effect.fn.Return<EffectCodexSchema.DynamicToolCallResponse> {
+  if (
+    payload.namespace !== T3_CODEX_DYNAMIC_TOOL_NAMESPACE ||
+    payload.tool !== T3_CODEX_WAIT_TOOL_NAME
+  ) {
+    return textResponse(
+      false,
+      `Unknown T3 Code dynamic tool: ${payload.namespace ?? "<none>"}.${payload.tool}`,
+    );
+  }
+
+  const decoded = yield* Effect.result(decodeT3CodexWaitArguments(payload.arguments));
+  if (Result.isFailure(decoded)) {
+    return textResponse(
+      false,
+      `Invalid wait arguments. durationMs must be an integer from ${T3_CODEX_WAIT_MIN_DURATION_MS} through ${T3_CODEX_WAIT_MAX_DURATION_MS}.`,
+    );
+  }
+
+  const outcome = yield* Effect.sleep(Duration.millis(decoded.success.durationMs)).pipe(
+    Effect.as("elapsed" as const),
+    Effect.raceFirst(cancelled.pipe(Effect.as("cancelled" as const))),
+  );
+
+  return outcome === "elapsed"
+    ? textResponse(true, `Wait completed after ${decoded.success.durationMs} ms.`)
+    : textResponse(false, "Wait cancelled because the turn was interrupted or the session closed.");
+});
+
+interface PendingWait {
+  readonly turnId: string;
+  readonly cancelled: Deferred.Deferred<void>;
+}
+
+interface WaitRegistryState {
+  readonly pending: Map<string, PendingWait>;
+  readonly cancelledTurnIds: Set<string>;
+  readonly closed: boolean;
+}
+
+export interface T3CodexDynamicToolWaitRegistry {
+  readonly handle: (
+    payload: EffectCodexSchema.DynamicToolCallParams,
+  ) => Effect.Effect<EffectCodexSchema.DynamicToolCallResponse>;
+  readonly cancelTurn: (turnId: string) => Effect.Effect<void>;
+  readonly cancelAll: Effect.Effect<void>;
+}
+
+export const makeT3CodexDynamicToolWaitRegistry = Effect.fn("makeT3CodexDynamicToolWaitRegistry")(
+  function* (): Effect.fn.Return<T3CodexDynamicToolWaitRegistry> {
+    const stateRef = yield* Ref.make<WaitRegistryState>({
+      pending: new Map(),
+      cancelledTurnIds: new Set(),
+      closed: false,
+    });
+
+    const completeCancellations = (pendingWaits: ReadonlyArray<PendingWait>) =>
+      Effect.forEach(
+        pendingWaits,
+        (pending) => Deferred.succeed(pending.cancelled, undefined).pipe(Effect.ignore),
+        { discard: true },
+      );
+
+    return {
+      handle: (payload) =>
+        Effect.gen(function* () {
+          const cancelled = yield* Deferred.make<void>();
+          const registered = yield* Ref.modify(stateRef, (current) => {
+            if (current.closed || current.cancelledTurnIds.has(payload.turnId)) {
+              return [false, current] as const;
+            }
+            const pending = new Map(current.pending);
+            pending.set(payload.callId, { turnId: payload.turnId, cancelled });
+            return [true, { ...current, pending }] as const;
+          });
+
+          return yield* handleT3CodexDynamicToolCall(
+            payload,
+            registered ? Deferred.await(cancelled) : Effect.void,
+          ).pipe(
+            Effect.ensuring(
+              Ref.update(stateRef, (current) => {
+                if (!current.pending.has(payload.callId)) {
+                  return current;
+                }
+                const pending = new Map(current.pending);
+                pending.delete(payload.callId);
+                return { ...current, pending };
+              }),
+            ),
+          );
+        }),
+      cancelTurn: (turnId) =>
+        Ref.modify(stateRef, (current) => {
+          const cancelledTurnIds = new Set(current.cancelledTurnIds);
+          cancelledTurnIds.add(turnId);
+          const pending = Array.from(current.pending.values()).filter(
+            (wait) => wait.turnId === turnId,
+          );
+          return [pending, { ...current, cancelledTurnIds }] as const;
+        }).pipe(Effect.flatMap(completeCancellations)),
+      cancelAll: Ref.modify(stateRef, (current) => [
+        Array.from(current.pending.values()),
+        { ...current, closed: true },
+      ]).pipe(Effect.flatMap(completeCancellations)),
+    };
+  },
+);