From b4277ab8db90fdb422efbcb2f64de9bd0ab5cde5 Mon Sep 17 00:00:00 2001 From: Yinjie Wang Date: Sat, 4 Apr 2026 01:44:04 -0500 Subject: [PATCH] Add files via upload --- extensions/rl-training-headers/README.md | 2 +- extensions/rl-training-headers/index.test.ts | 73 ++++++++++++++++++++ extensions/rl-training-headers/index.ts | 20 ++---- 3 files changed, 81 insertions(+), 14 deletions(-) create mode 100644 extensions/rl-training-headers/index.test.ts diff --git a/extensions/rl-training-headers/README.md b/extensions/rl-training-headers/README.md index fd54629..d26dac9 100644 --- a/extensions/rl-training-headers/README.md +++ b/extensions/rl-training-headers/README.md @@ -52,7 +52,7 @@ You can customize the header names in `~/.openclaw/openclaw.json` under `plugins The plugin hooks into the `before_prompt_build` lifecycle event to capture the current session ID and turn type (derived from the `trigger` field: `"user"` → `main`, `"heartbeat"` / `"memory"` / `"cron"` → `side`). -It then patches `globalThis.fetch` to inject these headers into all outgoing POST requests during an active agent run. Headers are cleared after the agent run completes (`agent_end`). +It then patches `globalThis.fetch` to inject these headers into outgoing POST requests during an active agent run. Header state is stored with async-local per-run context, so parallel sessions keep separate session IDs and turn types. ## Extracting training data diff --git a/extensions/rl-training-headers/index.test.ts b/extensions/rl-training-headers/index.test.ts new file mode 100644 index 0000000..666aa6e --- /dev/null +++ b/extensions/rl-training-headers/index.test.ts @@ -0,0 +1,73 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import register from "./index.js"; + +describe("rl-training-headers plugin", () => { + const hooks: Record = {}; + const originalFetch = globalThis.fetch; + + beforeEach(() => { + vi.clearAllMocks(); + for (const key of Object.keys(hooks)) { + delete hooks[key]; + } + }); + + afterEach(() => { + globalThis.fetch = originalFetch; + vi.restoreAllMocks(); + }); + + it("registers the before_prompt_build hook", () => { + const api = createApi(hooks); + + register(api as any); + + expect(api.on).toHaveBeenCalledWith("before_prompt_build", expect.any(Function)); + expect(api.logger.info).toHaveBeenCalledWith( + "rl-training-headers: activated (fetch patched)", + ); + }); + + it("keeps session headers isolated across concurrent runs", async () => { + const seenSessionIds: string[] = []; + const seenTurnTypes: string[] = []; + globalThis.fetch = vi.fn(async (_input: RequestInfo | URL, init?: RequestInit) => { + const headers = new Headers(init?.headers); + seenSessionIds.push(headers.get("X-Session-Id") ?? ""); + seenTurnTypes.push(headers.get("X-Turn-Type") ?? ""); + return new Response(null, { status: 200 }); + }) as typeof globalThis.fetch; + + register(createApi(hooks) as any); + + const runSession = async (sessionId: string, trigger: string) => { + await Promise.resolve(); + hooks.before_prompt_build?.({}, { sessionId, trigger }); + await Promise.resolve(); + await globalThis.fetch("https://example.test/llm", { method: "POST" }); + }; + + await Promise.all([ + runSession("session-a", "user"), + runSession("session-b", "heartbeat"), + ]); + + expect(seenSessionIds).toEqual(["session-a", "session-b"]); + expect(seenTurnTypes).toEqual(["main", "side"]); + }); +}); + +function createApi(hooks: Record) { + return { + pluginConfig: {}, + logger: { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + }, + on: vi.fn((name: string, handler: Function) => { + hooks[name] = handler; + }), + }; +} diff --git a/extensions/rl-training-headers/index.ts b/extensions/rl-training-headers/index.ts index bc6d8c1..f7f6bc1 100644 --- a/extensions/rl-training-headers/index.ts +++ b/extensions/rl-training-headers/index.ts @@ -1,3 +1,4 @@ +import { AsyncLocalStorage } from "node:async_hooks"; import type { OpenClawPluginApi } from "openclaw/plugin-sdk"; type RlTrainingConfig = { @@ -18,10 +19,7 @@ const SIDE_TRIGGERS = new Set(["heartbeat", "memory", "cron"]); export default function register(api: OpenClawPluginApi) { const config = resolveConfig(api); - - // Pending headers to inject into the next LLM fetch request. - // Set during before_prompt_build, consumed by the patched fetch, cleared on agent_end. - let pendingHeaders: Record | null = null; + const headerStore = new AsyncLocalStorage>(); const originalFetch = globalThis.fetch; @@ -29,10 +27,10 @@ export default function register(api: OpenClawPluginApi) { input: RequestInfo | URL, init?: RequestInit, ): Promise { - if (pendingHeaders && init?.method?.toUpperCase() === "POST") { - const extra = pendingHeaders; + const scopedHeaders = headerStore.getStore(); + if (scopedHeaders && init?.method?.toUpperCase() === "POST") { const merged = new Headers(init.headers); - for (const [k, v] of Object.entries(extra)) { + for (const [k, v] of Object.entries(scopedHeaders)) { // Plugin headers go first; per-request headers can still override. if (!merged.has(k)) { merged.set(k, v); @@ -46,16 +44,12 @@ export default function register(api: OpenClawPluginApi) { api.on("before_prompt_build", (_event, ctx) => { const sessionId = ctx.sessionId ?? ""; const turnType = SIDE_TRIGGERS.has(ctx.trigger ?? "") ? "side" : "main"; - pendingHeaders = { + headerStore.enterWith({ [config.sessionIdHeader]: sessionId, [config.turnTypeHeader]: turnType, - }; + }); return {}; }); - api.on("agent_end", () => { - pendingHeaders = null; - }); - api.logger.info("rl-training-headers: activated (fetch patched)"); }