chore: agent lock

Author: hugodutka

packages/blink/src/agent/client/index.ts

@@ -12,6 +12,7 @@ import { convertResponseToUIMessageStream } from "../internal/convert-response-t
 import type { ID } from "../types";
 import type { UIOptions, UIOptionsSchema } from "../ui";
 import { APIServerURLEnvironmentVariable } from "../constants";
+import { RWLock } from "../../local/rw-lock";
 
 export { APIServerURLEnvironmentVariable };
 
@@ -30,10 +31,12 @@ export type CapabilitiesResponse = Awaited<ReturnType<Client["capabilities"]>>;
 export class Client {
   public readonly baseUrl: string;
   private readonly client: ReturnType<typeof hc<typeof api>>;
+  public readonly agentLock: RWLock;
 
   public constructor(options: ClientOptions) {
     this.client = hc<typeof api>(options.baseUrl);
     this.baseUrl = options.baseUrl;
+    this.agentLock = new RWLock();
   }
 
   /**

packages/blink/src/local/chat-manager.test.ts

@@ -5,6 +5,7 @@ import { join } from "node:path";
 import type { UIMessage, UIMessageChunk } from "ai";
 import { ChatManager, type ChatStatus } from "./chat-manager";
 import { createDiskStore } from "./disk-store";
+import { RWLock } from "./rw-lock";
 import type { StoredChat, StoredMessage } from "./types";
 import type { Client } from "../agent/client";
 
@@ -14,6 +15,7 @@ function createMockAgent(
 ): Client & { chatCalls: any[] } {
   const chatCalls: any[] = [];
   return {
+    agentLock: new RWLock(),
     chatCalls,
     chat: async ({ messages, signal }: any) => {
       chatCalls.push({ messages, signal });
@@ -69,6 +71,7 @@ function createMockAgent(
 // Helper to create a slow-streaming agent (yields control between chunks)
 function createSlowAgent(chunks: number = 5): Client {
   return {
+    agentLock: new RWLock(),
     chat: async ({ signal }: any) => {
       const stream = new ReadableStream<UIMessageChunk>({
         async start(controller) {

packages/blink/src/local/chat-manager.ts

@@ -383,6 +383,8 @@ export class ChatManager {
           };
         });
 
+        // Acquire read lock on agent to prevent it from being disposed while streaming.
+        using _agentLock = await this.agent.agentLock.read();
         // Stream agent response
         const streamStartTime = performance.now();
         const stream = await runAgent({
@@ -541,7 +543,7 @@ export class ChatManager {
     this.disposed = true;
     this.watcher.dispose();
     this.listeners.clear();
-    this.abortController?.abort();
+    this.stopStreaming();
   }
 
   private resetChatState(): void {

packages/blink/src/react/use-agent.ts

@@ -7,6 +7,7 @@ import {
   type CapabilitiesResponse,
   APIServerURLEnvironmentVariable,
 } from "../agent/client";
+import { RWLock } from "../local/rw-lock";
 
 export interface AgentLog {
   readonly level: "log" | "error";
@@ -53,6 +54,8 @@ export default function useAgent(options: UseAgentOptions) {
     setAgent(undefined);
     setCapabilities(undefined);
 
+    let lock: RWLock | undefined;
+
     (async () => {
       const port = await getRandomPort();
       const proc = spawn("node", ["--no-deprecation", buildResult.entry], {
@@ -64,11 +67,27 @@ export default function useAgent(options: UseAgentOptions) {
           PORT: port.toString(),
           HOST: "127.0.0.1",
         },
+        // keep the child process tied to the parent process
+        detached: false,
       });
-      controller.signal.addEventListener("abort", () => {
+      const cleanup = () => {
         try {
           proc.kill();
         } catch {}
+      };
+
+      // Clean up - when the parent process exits, kill the child process.
+      process.once("exit", cleanup);
+      process.once("SIGINT", cleanup);
+      process.once("SIGTERM", cleanup);
+      process.once("uncaughtException", cleanup);
+
+      controller.signal.addEventListener("abort", () => {
+        process.off("exit", cleanup);
+        process.off("SIGINT", cleanup);
+        process.off("SIGTERM", cleanup);
+        process.off("uncaughtException", cleanup);
+        cleanup();
       });
       let ready = false;
       proc.stdout.on("data", (data) => {
@@ -114,6 +133,7 @@ export default function useAgent(options: UseAgentOptions) {
       const client = new Client({
         baseUrl: `http://127.0.0.1:${port}`,
       });
+      lock = client.agentLock;
       // Wait for the health endpoint to be alive.
       while (!controller.signal.aborted) {
         try {
@@ -139,7 +159,12 @@ export default function useAgent(options: UseAgentOptions) {
     });
     return () => {
       isCleanup = true;
-      controller.abort();
+      (async () => {
+        // Acquire write lock before cleaning up this agent instance
+        // This waits for any active streams using this agent to complete
+        using _writeLock = await lock?.write();
+        controller.abort();
+      })();
     };
   }, [buildResult, env, apiServerUrl]);
 

packages/blink/src/react/use-edit-agent.ts

@@ -6,6 +6,7 @@ import {
   getEditModeModel,
   type EditAgent,
 } from "../edit/agent";
+import { RWLock } from "../local/rw-lock";
 
 export interface UseEditAgentOptions {
   readonly directory: string;
@@ -35,6 +36,8 @@ export default function useEditAgent(options: UseEditAgentOptions) {
 
     setMissingApiKey(false);
 
+    let lock: RWLock | undefined;
+
     (async () => {
       // Create the edit agent
       editAgentRef.current = createEditAgent({
@@ -64,6 +67,7 @@ export default function useEditAgent(options: UseEditAgentOptions) {
       const editClient = new Client({
         baseUrl: `http://127.0.0.1:${port}`,
       });
+      lock = editClient.agentLock;
 
       // Wait for health check
       while (!controller.signal.aborted) {
@@ -88,7 +92,12 @@ export default function useEditAgent(options: UseEditAgentOptions) {
 
     return () => {
       isCleanup = true;
-      controller.abort();
+      (async () => {
+        // Acquire write lock before cleaning up this edit agent instance
+        // This waits for any active streams using this agent to complete
+        using _writeLock = await lock?.write();
+        controller.abort();
+      })();
     };
   }, [
     options.directory,