Adding hooks to record getlateststate for testing purposes (#6072)

Adding hooks to record getlateststate for testing purposes #6072

src/core/controller/grpc-handler.ts

@@ -1,10 +1,9 @@
+import { Controller } from "@core/controller/index"
 import { serviceHandlers } from "@generated/hosts/vscode/protobus-services"
+import { GrpcRecorderBuilder } from "@/core/controller/grpc-recorder/grpc-recorder.builder"
+import { GrpcRequestRegistry } from "@/core/controller/grpc-request-registry"
 import { ExtensionMessage } from "@/shared/ExtensionMessage"
 import { GrpcCancel, GrpcRequest } from "@/shared/WebviewMessage"
-import { GrpcRecorder, IRecorder } from "./grpc-recorder/grpc-recorder"
-import { LogFileHandler } from "./grpc-recorder/log-file-handler"
-import { GrpcRequestRegistry } from "./grpc-request-registry"
-import { Controller } from "./index"
 
 /**
  * Type definition for a streaming response handler
@@ -17,30 +16,17 @@ export type StreamingResponseHandler<TResponse> = (
 
 export type PostMessageToWebview = (message: ExtensionMessage) => Thenable<boolean | undefined>
 
-// Initialize the recorder as a singleton
-let recorder: IRecorder | null = null
-
-/**
- * Gets or creates the GrpcRecorder instance
- */
-function getRecorder(): IRecorder {
-	if (!recorder) {
-		recorder = GrpcRecorder.builder()
-			.enableIf(process.env.GRPC_RECORDER_ENABLED === "true" && process.env.CLINE_ENVIRONMENT === "local")
-			.withLogFileHandler(new LogFileHandler())
-			.build()
-	}
-	return recorder
-}
-
 /**
  * Creates a middleware wrapper for recording gRPC requests and responses
  */
-function withRecordingMiddleware(postMessage: PostMessageToWebview): PostMessageToWebview {
+function withRecordingMiddleware(postMessage: PostMessageToWebview, controller: Controller): PostMessageToWebview {
 	return async (response: ExtensionMessage) => {
 		if (response?.grpc_response) {
 			try {
-				getRecorder().recordResponse(response.grpc_response.request_id, response.grpc_response)
+				GrpcRecorderBuilder.getRecorder(controller).recordResponse(
+					response.grpc_response.request_id,
+					response.grpc_response,
+				)
 			} catch (e) {
 				console.warn("Failed to record gRPC response:", e)
 			}
@@ -52,9 +38,9 @@ function withRecordingMiddleware(postMessage: PostMessageToWebview): PostMessage
 /**
  * Records gRPC request with error handling
  */
-function recordRequest(request: GrpcRequest): void {
+function recordRequest(request: GrpcRequest, controller: Controller): void {
 	try {
-		getRecorder().recordRequest(request)
+		GrpcRecorderBuilder.getRecorder(controller).recordRequest(request)
 	} catch (e) {
 		console.warn("Failed to record gRPC request:", e)
 	}
@@ -68,10 +54,10 @@ export async function handleGrpcRequest(
 	postMessageToWebview: PostMessageToWebview,
 	request: GrpcRequest,
 ): Promise<void> {
-	recordRequest(request)
+	recordRequest(request, controller)
 
 	// Create recording middleware wrapper
-	const postMessageWithRecording = withRecordingMiddleware(postMessageToWebview)
+	const postMessageWithRecording = withRecordingMiddleware(postMessageToWebview, controller)
 
 	if (request.is_streaming) {
 		await handleStreamingRequest(controller, postMessageWithRecording, request)

src/core/controller/grpc-recorder/grpc-recorder.builder.ts

@@ -1,6 +1,8 @@
-import { GrpcRequestFilter } from "@core/controller/grpc-recorder/types"
+import { GrpcPostRecordHook, GrpcRequestFilter } from "@core/controller/grpc-recorder/types"
+import { Controller } from "@/core/controller"
 import { GrpcRecorder, GrpcRecorderNoops, IRecorder } from "@/core/controller/grpc-recorder/grpc-recorder"
 import { LogFileHandler, LogFileHandlerNoops } from "@/core/controller/grpc-recorder/log-file-handler"
+import { testHooks } from "@/core/controller/grpc-recorder/test-hooks"
 
 /**
  * A builder class for constructing a gRPC recorder instance.
@@ -13,6 +15,7 @@ export class GrpcRecorderBuilder {
 	private fileHandler: LogFileHandler | null = null
 	private enabled: boolean = true
 	private filters: GrpcRequestFilter[] = []
+	private hooks: GrpcPostRecordHook[] = []
 
 	public withLogFileHandler(handler: LogFileHandler): this {
 		this.fileHandler = handler
@@ -29,7 +32,28 @@ export class GrpcRecorderBuilder {
 		return this
 	}
 
-	public build(): IRecorder {
+	public withPostRecordHooks(...hooks: GrpcPostRecordHook[]): this {
+		this.hooks.push(...hooks)
+		return this
+	}
+
+	// Initialize the recorder as a singleton
+	private static recorder: IRecorder
+
+	/**
+	 * Gets or creates the GrpcRecorder instance
+	 */
+	static getRecorder(controller: Controller): IRecorder {
+		if (!GrpcRecorderBuilder.recorder) {
+			GrpcRecorderBuilder.recorder = GrpcRecorder.builder()
+				.enableIf(process.env.GRPC_RECORDER_ENABLED === "true" && process.env.CLINE_ENVIRONMENT === "local")
+				.withLogFileHandler(new LogFileHandler())
+				.build(controller)
+		}
+		return GrpcRecorderBuilder.recorder
+	}
+
+	public build(controller?: Controller): IRecorder {
 		if (!this.enabled) {
 			return new GrpcRecorderNoops()
 		}
@@ -39,8 +63,13 @@ export class GrpcRecorderBuilder {
 			filters = filters.concat(this.filters)
 		}
 
+		let hooks: GrpcPostRecordHook[] = hooksFromEnv(controller)
+		if (this.hooks.length > 0) {
+			hooks = hooks.concat(this.hooks)
+		}
+
 		const handler = this.fileHandler ?? new LogFileHandlerNoops()
-		return new GrpcRecorder(handler, filters)
+		return new GrpcRecorder(handler, filters, hooks)
 	}
 }
 
@@ -61,3 +90,13 @@ function testFilters(): GrpcRequestFilter[] {
 	 */
 	return [(req) => req.is_streaming, (req) => ["cline.UiService", "cline.McpService", "cline.WebService"].includes(req.service)]
 }
+
+function hooksFromEnv(controller?: Controller): GrpcPostRecordHook[] {
+	const hooks: GrpcPostRecordHook[] = []
+
+	if (controller && process.env.GRPC_RECORDER_TESTS_FILTERS_ENABLED === "true") {
+		hooks.push(...testHooks(controller))
+	}
+
+	return hooks
+}

src/core/controller/grpc-recorder/grpc-recorder.test.ts

@@ -113,7 +113,6 @@ describe("grpc-recorder", () => {
 				is_streaming: false,
 			})
 
-			// we expect to filter out this unwanted request
 			expect(sessionLog.entries).length(3)
 		})
 
@@ -147,5 +146,69 @@ describe("grpc-recorder", () => {
 			sessionLog = customRecorder.getSessionLog()
 			expect(sessionLog.entries).length(0)
 		})
+
+		it("cleanupSyntheticEntries removes synthetic entries from session log", async () => {
+			const testRecorder = GrpcRecorder.builder().enableIf(true).build()
+
+			// Add regular request
+			testRecorder.recordRequest({
+				service: "regular-service",
+				method: "regular-method",
+				message: "regular-message",
+				request_id: "regular-id",
+				is_streaming: false,
+			})
+
+			// Add synthetic request
+			testRecorder.recordRequest(
+				{
+					service: "synthetic-service",
+					method: "synthetic-method",
+					message: "synthetic-message",
+					request_id: "synthetic-id",
+					is_streaming: false,
+				},
+				true, // synthetic = true
+			)
+
+			let sessionLog = testRecorder.getSessionLog()
+			expect(sessionLog.entries).length(2)
+
+			testRecorder.cleanupSyntheticEntries()
+
+			sessionLog = testRecorder.getSessionLog()
+			expect(sessionLog.entries).length(1)
+			expect(sessionLog.entries[0].requestId).equal("regular-id")
+		})
+
+		it("recordResponse executes post-record hooks", async () => {
+			let hookExecuted = false
+			let hookEntry: any = null
+
+			const mockHook = async (entry: any) => {
+				hookExecuted = true
+				hookEntry = entry
+			}
+
+			const testRecorder = GrpcRecorder.builder().withPostRecordHooks(mockHook).enableIf(true).build()
+
+			testRecorder.recordRequest({
+				service: "test-service",
+				method: "test-method",
+				message: "test-message",
+				request_id: "test-id",
+				is_streaming: false,
+			})
+
+			testRecorder.recordResponse("test-id", {
+				request_id: "test-id",
+				message: "response-message",
+				error: "",
+			})
+
+			expect(hookExecuted).to.be.true
+			expect(hookEntry).to.not.be.null
+			expect(hookEntry.requestId).equal("test-id")
+		})
 	})
 })

src/core/controller/grpc-recorder/grpc-recorder.ts

@@ -2,7 +2,13 @@ import { GrpcResponse } from "@shared/ExtensionMessage"
 import { GrpcRequest } from "@shared/WebviewMessage"
 import { GrpcRecorderBuilder } from "@/core/controller/grpc-recorder/grpc-recorder.builder"
 import { ILogFileHandler } from "@/core/controller/grpc-recorder/log-file-handler"
-import { GrpcLogEntry, GrpcRequestFilter, GrpcSessionLog, SessionStats } from "@/core/controller/grpc-recorder/types"
+import {
+	GrpcLogEntry,
+	GrpcPostRecordHook,
+	GrpcRequestFilter,
+	GrpcSessionLog,
+	SessionStats,
+} from "@/core/controller/grpc-recorder/types"
 
 export class GrpcRecorderNoops implements IRecorder {
 	recordRequest(_request: GrpcRequest): void {}
@@ -14,13 +20,15 @@ export class GrpcRecorderNoops implements IRecorder {
 			entries: [],
 		}
 	}
+	cleanupSyntheticEntries(): void {}
 }
 
 export interface IRecorder {
-	recordRequest(request: GrpcRequest): void
+	recordRequest(request: GrpcRequest, synthetic?: boolean): void
 	recordResponse(requestId: string, response: GrpcResponse): void
 	recordError(requestId: string, error: string): void
 	getSessionLog(): GrpcSessionLog
+	cleanupSyntheticEntries(): void
 }
 
 /**
@@ -39,6 +47,7 @@ export class GrpcRecorder implements IRecorder {
 	constructor(
 		private fileHandler: ILogFileHandler,
 		private requestFilters: GrpcRequestFilter[] = [],
+		private postRecordHooks: GrpcPostRecordHook[] = [],
 	) {
 		this.sessionLog = {
 			startTime: new Date().toISOString(),
@@ -63,7 +72,7 @@ export class GrpcRecorder implements IRecorder {
 	 *
 	 * @param request - The incoming gRPC request.
 	 */
-	public recordRequest(request: GrpcRequest): void {
+	public recordRequest(request: GrpcRequest, synthetic: boolean = false): void {
 		if (this.shouldFilter(request)) {
 			return
 		}
@@ -77,6 +86,7 @@ export class GrpcRecorder implements IRecorder {
 				message: request.message,
 			},
 			status: "pending",
+			meta: { synthetic },
 		}
 
 		this.pendingRequests.set(request.request_id, {
@@ -106,6 +116,7 @@ export class GrpcRecorder implements IRecorder {
 	 */
 	public recordResponse(requestId: string, response: GrpcResponse): void {
 		const pendingRequest = this.pendingRequests.get(requestId)
+
 		if (!pendingRequest) {
 			console.warn(`No pending request found for response with ID: ${requestId}`)
 			return
@@ -130,6 +141,30 @@ export class GrpcRecorder implements IRecorder {
 		this.sessionLog.stats = this.getStats()
 
 		this.flushLogAsync()
+
+		this.runHooks(entry).catch((e) => console.error("Post-record hook failed:", e))
+	}
+
+	private async runHooks(entry: GrpcLogEntry): Promise<void> {
+		if (entry.meta?.synthetic) return
+		for (const hook of this.postRecordHooks) {
+			await hook(entry)
+		}
+	}
+
+	public cleanupSyntheticEntries(): void {
+		// Remove synthetic entries from session log
+		this.sessionLog.entries = this.sessionLog.entries.filter((entry) => !entry.meta?.synthetic)
+
+		// clean up from pending requests if needed
+		for (const [requestId, pendingRequest] of this.pendingRequests.entries()) {
+			if (pendingRequest.entry.meta?.synthetic) {
+				this.pendingRequests.delete(requestId)
+			}
+		}
+
+		this.sessionLog.stats = this.getStats()
+		this.flushLogAsync()
 	}
 
 	/**

src/core/controller/grpc-recorder/test-hooks.test.ts

@@ -0,0 +1,69 @@
+import { afterEach, beforeEach, describe, it } from "mocha"
+import "should"
+import * as sinon from "sinon"
+import { Controller } from ".."
+import { IRecorder } from "./grpc-recorder"
+import { GrpcRecorderBuilder } from "./grpc-recorder.builder"
+import { testHooks } from "./test-hooks"
+import { GrpcLogEntry } from "./types"
+
+describe("test-hooks", () => {
+	let cleanupSyntheticEntriesStub: sinon.SinonStub
+	let recordRequestStub: sinon.SinonStub
+	let recordResponseStub: sinon.SinonStub
+	let getRecorderStub: sinon.SinonStub
+
+	beforeEach(() => {
+		cleanupSyntheticEntriesStub = sinon.stub()
+		recordRequestStub = sinon.stub()
+		recordResponseStub = sinon.stub()
+
+		const mockRecorder: IRecorder = {
+			cleanupSyntheticEntries: cleanupSyntheticEntriesStub,
+			recordRequest: recordRequestStub,
+			recordResponse: recordResponseStub,
+			recordError: sinon.stub(),
+			getSessionLog: sinon.stub().returns({ startTime: "", entries: [] }),
+		}
+
+		getRecorderStub = sinon.stub(GrpcRecorderBuilder, "getRecorder").returns(mockRecorder)
+	})
+
+	afterEach(() => {
+		sinon.restore()
+	})
+
+	it("should return an array of post-record hooks", () => {
+		const mockController = {} as Controller
+		const hooks = testHooks(mockController)
+
+		hooks.should.be.an.Array()
+		hooks.should.have.length(1)
+		hooks[0].should.be.a.Function()
+	})
+
+	it("should execute hook and call recorder methods", async () => {
+		const mockController = {
+			getStateToPostToWebview: sinon.stub().returns({}),
+		} as any as Controller
+
+		const hooks = testHooks(mockController)
+
+		const mockEntry: GrpcLogEntry = {
+			requestId: "test-request-id",
+			service: "TestService",
+			method: "testMethod",
+			isStreaming: false,
+			request: { message: {} },
+			status: "pending",
+		}
+
+		await hooks[0](mockEntry)
+
+		// Validate sinon stub calls
+		sinon.assert.calledWith(getRecorderStub, mockController)
+		sinon.assert.calledOnce(cleanupSyntheticEntriesStub)
+		sinon.assert.calledOnce(recordRequestStub)
+		sinon.assert.calledOnce(recordResponseStub)
+	})
+})

src/core/controller/grpc-recorder/test-hooks.ts

@@ -0,0 +1,38 @@
+import { Controller } from "@/core/controller"
+import { GrpcRecorderBuilder } from "@/core/controller/grpc-recorder/grpc-recorder.builder"
+import { GrpcPostRecordHook } from "@/core/controller/grpc-recorder/types"
+import { getLatestState } from "@/core/controller/state/getLatestState"
+
+// Add 50ms delay by default to ensure we get the latest state
+const TEST_HOOK_LATEST_STATE_DELAY = 50
+
+export function testHooks(controller: Controller): GrpcPostRecordHook[] {
+	return [
+		async (entry) => {
+			GrpcRecorderBuilder.getRecorder(controller).cleanupSyntheticEntries()
+
+			await new Promise((resolve) => setTimeout(resolve, TEST_HOOK_LATEST_STATE_DELAY))
+
+			const requestId = entry.requestId
+
+			// Record synthetic "getLatestState" request
+			GrpcRecorderBuilder.getRecorder(controller).recordRequest(
+				{
+					service: "cline.StateService",
+					method: "getLatestState",
+					message: {},
+					request_id: requestId,
+					is_streaming: false,
+				},
+				true,
+			)
+
+			const state = await getLatestState(controller, {})
+
+			GrpcRecorderBuilder.getRecorder(controller).recordResponse(requestId, {
+				request_id: requestId,
+				message: state,
+			})
+		},
+	]
+}

src/core/controller/grpc-recorder/types.ts

@@ -1,5 +1,7 @@
 import { GrpcRequest } from "@/shared/WebviewMessage"
 
+export type GrpcPostRecordHook = (entry: GrpcLogEntry, controller?: any) => Promise<void> | void
+
 export type GrpcRequestFilter = (request: GrpcRequest) => boolean
 
 export interface GrpcLogEntry {
@@ -18,6 +20,7 @@ export interface GrpcLogEntry {
 	}
 	duration?: number
 	status: "pending" | "completed" | "error"
+	meta?: { synthetic?: boolean }
 }
 
 export interface SessionStats {