Repository: stateful/vscode-runme
Pull request #281 (2023-03-06 22:29:00)

Auto-Restart server on `onClose`
Makes runme server resistant against closing, as long as it is caught by the `onClose` callback.

Might want to think about polling/regular health checks in case something more dubious occurs. (Long-term)

Should be functional, just need to write/fix tests at this point.

sourishkrout: K, I'm going to test this end-to-end shortly.
sourishkrout: Worked on Mac and Windows (both WSL and regular).
sourishkrout: This PR fixes https://github.com/stateful/vscode-runme/issues/225.

Patch:
diff --git a/__mocks__/vscode.ts b/__mocks__/vscode.ts
index 6071dc64..da30326c 100644
--- a/__mocks__/vscode.ts
+++ b/__mocks__/vscode.ts
@@ -4,6 +4,8 @@ import { join } from 'node:path'
 
 import { vi } from 'vitest'
 import { URI } from 'vscode-uri'
+import type { Disposable } from 'vscode'
+
 
 export const notebooks = {
   createNotebookController: vi.fn().mockReturnValue({
@@ -116,7 +118,6 @@ export enum NotebookCellKind {
 }
 
 export const TreeItem = vi.fn()
-export class EventEmitter {}
 export const Event = vi.fn()
 
 export enum TreeItemCollapsibleState {
@@ -193,3 +194,26 @@ export class NotebookCellOutput {
 export const ProgressLocation = {
   Window: 1
 }
+
+type MessageCallback<T> = (message: T) => void
+type Event<T> = (listener: MessageCallback<T>) => Disposable
+
+export class EventEmitter<T> {
+  listeners: MessageCallback<T>[] = []
+
+  event: Event<T> = (listener) => {
+    this.listeners.push(listener)
+
+    return {
+      dispose: () => {
+        this.listeners = this.listeners.filter(x => x !== listener)
+      }
+    }
+  }
+
+  fire(data: T) {
+    this.listeners.forEach(l => l(data))
+  }
+
+  dispose() { }
+}
diff --git a/src/extension/extension.ts b/src/extension/extension.ts
index 1cef3995..8f4347f3 100644
--- a/src/extension/extension.ts
+++ b/src/extension/extension.ts
@@ -31,22 +31,19 @@ export class RunmeExtension {
     const kernel = new Kernel(context)
     const grpcSerializer = kernel.hasExperimentEnabled('grpcSerializer')
     const grpcServer = kernel.hasExperimentEnabled('grpcServer')
+    const grpcRunner = kernel.hasExperimentEnabled('grpcRunner')
     const server = new RunmeServer(context.extensionUri, {
       retryOnFailure: true,
       maxNumberOfIntents: 2,
-    })
+    }, !grpcServer, grpcRunner)
+    grpcRunner && kernel.useRunner(new GrpcRunner(server))
+    const serializer = grpcSerializer ? new GrpcSerializer(context, server) : new WasmSerializer(context)
+
     /**
      * Start the Runme server
      */
-
     try {
-      if (grpcServer) {
-        await server.launch()
-        server.events.on('closed', () => {
-          window.showErrorMessage(`Runme server is not longer running,
-          please reload the window to start the server again`)
-        })
-      }
+      await server.launch()
     } catch (e) {
       // Unrecoverable error happened
       if (e instanceof RunmeServerError) {
@@ -57,10 +54,8 @@ export class RunmeExtension {
       return window.showErrorMessage('Failed to start Runme server, please try to reload the window')
     }
 
-    const serializer = grpcSerializer ? new GrpcSerializer(context) : new WasmSerializer(context)
     const treeViewer = new RunmeLauncherProvider(getDefaultWorkspace())
     const uriHandler = new RunmeUriHandler(context)
-    kernel.hasExperimentEnabled('grpcRunner') && kernel.useRunner(new GrpcRunner())
 
     context.subscriptions.push(
       kernel,
diff --git a/src/extension/grpc/client.ts b/src/extension/grpc/client.ts
index a9c187d6..a1622872 100644
--- a/src/extension/grpc/client.ts
+++ b/src/extension/grpc/client.ts
@@ -2,17 +2,9 @@
 import { ParserServiceClient } from '@buf/stateful_runme.community_timostamm-protobuf-ts/runme/parser/v1/parser_pb.client'
 // eslint-disable-next-line max-len
 import { RunnerServiceClient } from '@buf/stateful_runme.community_timostamm-protobuf-ts/runme/runner/v1/runner_pb.client'
-import { ChannelCredentials } from '@grpc/grpc-js'
 import { GrpcTransport } from '@protobuf-ts/grpc-transport'
 
-import { getGrpcHost } from '../utils'
-
-function initParserClient(): ParserServiceClient {
-  const transport = new GrpcTransport({
-    host: getGrpcHost(),
-    channelCredentials: ChannelCredentials.createInsecure(),
-  })
-
+function initParserClient(transport: GrpcTransport): ParserServiceClient {
   return new ParserServiceClient(transport)
 }
 
diff --git a/src/extension/kernel.ts b/src/extension/kernel.ts
index aa95fdd1..88ab7595 100644
--- a/src/extension/kernel.ts
+++ b/src/extension/kernel.ts
@@ -48,6 +48,7 @@ export class Kernel implements Disposable {
 
   protected runner?: IRunner
   protected environment?: IRunnerEnvironment
+  protected runnerReadyListener?: Disposable
 
   constructor(
     protected context: ExtensionContext,
@@ -80,6 +81,7 @@ export class Kernel implements Disposable {
     resetEnv()
     this.#controller.dispose()
     this.#disposables.forEach((d) => d.dispose())
+    this.runnerReadyListener?.dispose()
   }
 
   async #handleSaveNotebook({ uri, isUntitled, notebookType }: NotebookDocument) {
@@ -316,15 +318,23 @@ export class Kernel implements Disposable {
   }
 
   useRunner(runner: IRunner) {
+    this.runnerReadyListener?.dispose()
+
     if(this.#experiments.get('grpcRunner') && runner) {
       this.runner = runner
 
-      this.#disposables.push(runner)
-      runner.createEnvironment(
-        // copy env from process naively for now
-        // later we might want a more sophisticated approach/to bring this serverside
-        Object.entries(process.env).map(([k, v]) => `${k}=${v || ''}`)
-      ).then(env => { this.environment = env })
+      this.runnerReadyListener = runner.onReady(() => {
+        this.environment = undefined
+
+        runner.createEnvironment(
+          // copy env from process naively for now
+          // later we might want a more sophisticated approach/to bring this serverside
+          Object.entries(process.env).map(([k, v]) => `${k}=${v || ''}`)
+        ).then(env => {
+          if(this.runner !== runner) { return }
+          this.environment = env
+        })
+      })
     }
   }
 
diff --git a/src/extension/runner.ts b/src/extension/runner.ts
index 6fdf267b..a96407cc 100644
--- a/src/extension/runner.ts
+++ b/src/extension/runner.ts
@@ -1,4 +1,3 @@
-import { ChannelCredentials } from '@grpc/grpc-js'
 import { GrpcTransport } from '@protobuf-ts/grpc-transport'
 import { DuplexStreamingCall } from '@protobuf-ts/runtime-rpc/build/types/duplex-streaming-call'
 import {
@@ -21,8 +20,8 @@ import {
   Session,
 } from './grpc/runnerTypes'
 import { RunnerServiceClient } from './grpc/client'
-import { getGrpcHost } from './utils'
 import { getShellPath } from './executors/utils'
+import RunmeServer from './server/runmeServer'
 
 type ExecuteDuplex = DuplexStreamingCall<ExecuteRequest, ExecuteResponse>
 
@@ -46,6 +45,13 @@ export interface RunProgramOptions {
 }
 
 export interface IRunner extends Disposable {
+  /**
+   * Called when underlying transport is ready
+   *
+   * May be called multiple times if server restarts
+   */
+  readonly onReady: Event<void>
+
   close(): void
 
   createEnvironment(
@@ -65,9 +71,11 @@ export interface IRunner extends Disposable {
   ): Promise<boolean>
 }
 
-export interface IRunnerEnvironment extends DisposableAsync { }
+interface IRunnerChild extends DisposableAsync { }
 
-export interface IRunnerProgramSession extends DisposableAsync, Pseudoterminal {
+export interface IRunnerEnvironment extends IRunnerChild { }
+
+export interface IRunnerProgramSession extends IRunnerChild, Pseudoterminal {
   /**
    * Called when an unrecoverable error occurs, for instance a failure over gRPC
    * transport
@@ -98,23 +106,46 @@ export interface IRunnerProgramSession extends DisposableAsync, Pseudoterminal {
 }
 
 export class GrpcRunner implements IRunner {
-  protected readonly client: RunnerServiceClient
-  protected transport: GrpcTransport
+  protected client?: RunnerServiceClient
 
-  private children: WeakRef<DisposableAsync>[] = []
+  private children: WeakRef<IRunnerChild>[] = []
+  private disposables: Disposable[] = []
 
-  constructor() {
-    this.transport = new GrpcTransport({
-      host: getGrpcHost(),
-      channelCredentials: ChannelCredentials.createInsecure(),
-    })
+  protected _onReady = this.register(new EventEmitter<void>())
+  onReady = this._onReady.event
 
-    this.client = new RunnerServiceClient(this.transport)
+  constructor(protected server: RunmeServer) {
+    this.disposables.push(
+      server.onTransportReady(({ transport }) => this.initRunnerClient(transport))
+    )
+
+    this.disposables.push(
+      server.onClose(() => this.deinitRunnerClient())
+    )
+  }
+
+  private deinitRunnerClient() {
+    this.disposeChildren()
+    this.client = undefined
+  }
+
+  private initRunnerClient(transport?: GrpcTransport) {
+    this.deinitRunnerClient()
+    this.client = new RunnerServiceClient(transport ?? this.server.transport())
+    this._onReady.fire()
+  }
+
+  protected static assertClient(client: RunnerServiceClient|undefined): asserts client {
+    if(!client) {
+      throw new Error('Client is not active!')
+    }
   }
 
   async createProgramSession(
     opts: RunProgramOptions
   ): Promise<IRunnerProgramSession> {
+    GrpcRunner.assertClient(this.client)
+
     const session = new GrpcRunnerProgramSession(
       this.client,
       opts
@@ -125,16 +156,20 @@ export class GrpcRunner implements IRunner {
     return session
   }
 
-  createEnvironment(
+  async createEnvironment(
     envs?: string[],
     metadata?: { [index: string]: string }
   ) {
+    GrpcRunner.assertClient(this.client)
+
     const request = CreateSessionRequest.create({
       metadata, envs
     })
 
     try {
-      return this.client
+      const client = this.client
+
+      return client
         .createSession(request)
         .then(({ response: { session } }) => {
           if(!session) {
@@ -142,7 +177,7 @@ export class GrpcRunner implements IRunner {
           }
 
           const environment = new GrpcRunnerEnvironment(
-            this.client,
+            client,
             session
           )
 
@@ -162,6 +197,8 @@ export class GrpcRunner implements IRunner {
   async getEnvironmentVariables(
     environment: IRunnerEnvironment,
   ): Promise<Record<string, string> | undefined> {
+    GrpcRunner.assertClient(this.client)
+
     if(!(environment instanceof GrpcRunnerEnvironment)) {
       throw new Error('Invalid environment!')
     }
@@ -212,14 +249,17 @@ export class GrpcRunner implements IRunner {
     })
   }
 
-  close(): void {
-    this.transport.close()
-  }
+  close(): void { }
 
   async dispose(): Promise<void> {
+    this.disposables.forEach(d => d.dispose())
+    await this.disposeChildren().finally(() => this.close())
+  }
+
+  async disposeChildren(): Promise<void> {
     await Promise.all(
       this.children.map(c => c.deref()?.dispose())
-    ).finally(() => this.close())
+    )
   }
 
   /**
@@ -229,6 +269,11 @@ export class GrpcRunner implements IRunner {
   protected registerChild(d: DisposableAsync) {
     this.children.push(new WeakRef(d))
   }
+
+  protected register<T extends Disposable>(d: T): T {
+    this.disposables.push(d)
+    return d
+  }
 }
 
 export class GrpcRunnerProgramSession implements IRunnerProgramSession {
@@ -280,6 +325,9 @@ export class GrpcRunnerProgramSession implements IRunnerProgramSession {
       })
     )
 
+    this.register( this._onDidClose.event(() => this.dispose()) )
+    this.register( this._onInternalErr.event(() => this.dispose()) )
+
     this.session.responses.onMessage(({ stderrData, stdoutData, exitCode }) => {
       if(stdoutData.length > 0) {
         this._onStdoutRaw.fire(stdoutData)
diff --git a/src/extension/serializer.ts b/src/extension/serializer.ts
index b7858ab2..dbb49184 100644
--- a/src/extension/serializer.ts
+++ b/src/extension/serializer.ts
@@ -15,6 +15,7 @@ import {
   Disposable,
 } from 'vscode'
 import { v4 as uuidv4 } from 'uuid'
+import { GrpcTransport } from '@protobuf-ts/grpc-transport'
 
 import { Serializer } from '../types'
 
@@ -23,6 +24,7 @@ import { initParserClient, ParserServiceClient } from './grpc/client'
 import Languages from './languages'
 import { PLATFORM_OS } from './constants'
 import { canEditFile, initWasm } from './utils'
+import RunmeServer from './server/runmeServer'
 
 declare var globalThis: any
 const DEFAULT_LANG_ID = 'text'
@@ -288,13 +290,26 @@ export class GrpcSerializer extends SerializerBase {
   private client?: ParserServiceClient
   protected ready: ReadyPromise
 
-  constructor(protected context: ExtensionContext) {
+  private serverReadyListener: Disposable|undefined
+
+  constructor(
+    protected context: ExtensionContext,
+    protected server: RunmeServer,
+  ) {
     super(context)
 
     this.ready = new Promise((resolve) => {
-      this.client = initParserClient()
-      resolve()
+      const disposable = server.onTransportReady(() => {
+        disposable.dispose()
+        resolve()
+      })
     })
+
+    this.serverReadyListener = server.onTransportReady(({ transport }) => this.initParserClient(transport))
+  }
+
+  private initParserClient(transport?: GrpcTransport) {
+    this.client = initParserClient(transport ?? this.server.transport())
   }
 
   protected async saveNotebook(
@@ -332,4 +347,9 @@ export class GrpcSerializer extends SerializerBase {
     // we can remove ugly casting once we switch to GRPC
     return (notebook as unknown) as Serializer.Notebook
   }
+
+  public dispose(): void {
+    this.serverReadyListener?.dispose()
+    super.dispose()
+  }
 }
diff --git a/src/extension/server/runmeServer.ts b/src/extension/server/runmeServer.ts
index 712355c1..b58529f0 100644
--- a/src/extension/server/runmeServer.ts
+++ b/src/extension/server/runmeServer.ts
@@ -1,8 +1,9 @@
 import { ChildProcessWithoutNullStreams, spawn } from 'node:child_process'
 import fs from 'node:fs/promises'
-import EventEmitter from 'node:events'
 
-import { Disposable, Uri } from 'vscode'
+import { ChannelCredentials } from '@grpc/grpc-js'
+import { GrpcTransport } from '@protobuf-ts/grpc-transport'
+import { Disposable, Uri, EventEmitter } from 'vscode'
 
 import { SERVER_ADDRESS } from '../../constants'
 import { enableServerLogs, getBinaryPath, getPortNumber } from '../../utils/configuration'
@@ -23,7 +24,6 @@ export interface IServerConfig {
 }
 
 class RunmeServer implements Disposable {
-
     #port: number
     #process: ChildProcessWithoutNullStreams | undefined
     #binaryPath: Uri
@@ -33,9 +33,21 @@ class RunmeServer implements Disposable {
     #intent: number
     #acceptsIntents: number
     #acceptsInterval: number
-    events: EventEmitter
+    #disposables: Disposable[] = []
+    #transport?: GrpcTransport
+
+    readonly #onClose = this.register(new EventEmitter<{ code: number|null }>())
+    readonly #onTransportReady = this.register(new EventEmitter<{ transport: GrpcTransport }>())
+
+    readonly onClose = this.#onClose.event
+    readonly onTransportReady = this.#onTransportReady.event
 
-    constructor(extBasePath: Uri, options: IServerConfig) {
+    constructor(
+      extBasePath: Uri,
+      options: IServerConfig,
+      protected readonly externalServer: boolean,
+      protected readonly enableRunner = false
+    ) {
       this.#port = getPortNumber()
       this.#loggingEnabled = enableServerLogs()
       this.#binaryPath = getBinaryPath(extBasePath, process.platform)
@@ -44,17 +56,26 @@ class RunmeServer implements Disposable {
       this.#intent = 0
       this.#acceptsIntents = options.acceptsConnection?.intents || 50
       this.#acceptsInterval = options.acceptsConnection?.interval || 200
-      this.events = new EventEmitter()
     }
 
     dispose() {
-        this.events.removeAllListeners()
-        this.#process?.removeAllListeners()
-        this.#process?.kill()
+      this.#disposables.forEach(d => d.dispose())
+      this.disposeProcess()
+    }
+
+    private disposeProcess(process?: ChildProcessWithoutNullStreams) {
+      process ??= this.#process
+
+      if(process === this.#process) {
+        this.#process = undefined
+      }
+
+      process?.removeAllListeners()
+      process?.kill()
     }
 
     async isRunning(): Promise<boolean> {
-      const client = initParserClient()
+      const client = initParserClient(this.transport())
       try {
         const deserialRequest = DeserializeRequest.create({ source: Buffer.from('## Server running', 'utf-8') })
         const request = client.deserialize(deserialRequest)
@@ -68,11 +89,31 @@ class RunmeServer implements Disposable {
       }
     }
 
-    private address() {
+    protected address() {
       return `${SERVER_ADDRESS}:${this.#port}`
     }
 
-    async start(): Promise<string | RunmeServerError> {
+    protected channelCredentials() {
+      return ChannelCredentials.createInsecure()
+    }
+
+    protected closeTransport() {
+      this.#transport?.close()
+      this.#transport = undefined
+    }
+
+    transport() {
+      if(this.#transport) { return this.#transport }
+
+      this.#transport = new GrpcTransport({
+        host: this.address(),
+        channelCredentials: this.channelCredentials(),
+      })
+
+      return this.#transport
+    }
+
+    async start(): Promise<string> {
         const binaryLocation = this.#binaryPath.fsPath
 
         const binaryExists = await fs.access(binaryLocation)
@@ -87,46 +128,64 @@ class RunmeServer implements Disposable {
             throw new RunmeServerError('Cannot find server binary file')
         }
 
+        this.#port = getPortNumber()
         while (!(await isPortAvailable(this.#port))) {
           this.#port++
         }
 
-        this.#process = spawn(binaryLocation, [
-            'server',
-            '--address',
-            this.address()
-        ])
+        const args = [
+          'server',
+          '--address',
+          this.address()
+        ]
+
+        if(this.enableRunner) {
+          args.push('--runner')
+        }
 
-        this.#process.on('close', () => {
+        const process = spawn(binaryLocation, args)
+
+        process.on('close', (code) => {
             if (this.#loggingEnabled) {
-                console.log(`[Runme] Server process #${this.#process?.pid} closed`)
+                console.log(`[Runme] Server process #${this.#process?.pid} closed with code ${code}`)
             }
-            this.events.emit('closed')
+            this.#onClose.fire({ code })
+
+            this.disposeProcess(process)
+
+            // try to relaunch
+            this.launch()
         })
 
 
-        this.#process.stderr.once('data', () => {
+        process.stderr.once('data', () => {
             console.log(`[Runme] Server process #${this.#process?.pid} started on port ${this.#port}`)
         })
 
-        this.#process.stderr.on('data', (data) => {
+        process.stderr.on('data', (data) => {
             if (this.#loggingEnabled) {
                 console.log(data.toString())
             }
         })
 
+        this.#process = process
+
         return new Promise((resolve, reject) => {
-            this.#process!.stderr.on('data', (data) => {
-                const msg = data.toString()
-                try {
-                    const log = JSON.parse(msg)
-                    if (log.addr === this.address()) {
-                        return resolve(log.addr)
-                    }
-                } catch (err: any) {
-                    reject(new RunmeServerError(`Server failed, reason: ${msg || (err as Error).message}`))
-                }
-            })
+          const cb = (data: any) => {
+              const msg = data.toString()
+              try {
+                  const log = JSON.parse(msg)
+                  if (log.addr === this.address()) {
+                    return resolve(log.addr)
+                  }
+              } catch (err: any) {
+                  reject(new RunmeServerError(`Server failed, reason: ${msg || (err as Error).message}`))
+              } finally {
+                process.stderr.off('data', cb)
+              }
+          }
+
+          process.stderr.on('data', cb)
         })
     }
 
@@ -157,24 +216,51 @@ class RunmeServer implements Disposable {
         })
     }
 
-    async launch(): Promise<string | RunmeServerError> {
-        let addr
-        try {
-            addr = await this.start()
-        } catch (e) {
-            if (this.#retryOnFailure && this.#maxNumberOfIntents > this.#intent) {
-                this.#intent++
-                return this.launch()
-            }
-            throw new RunmeServerError(`Cannot start server. Error: ${(e as Error).message}`)
-        }
-        await this.acceptsConnection()
-        return addr
+    /**
+     * Tries to launch server, retrying if needed
+     *
+     * If `externalServer` is set, then this only attempts to connect to the
+     * server address
+     *
+     * @returns Address of server or error
+     */
+    async launch(): Promise<string> {
+      if (this.externalServer) {
+        await this.connect()
+        return this.address()
+      }
+
+      let addr
+      try {
+          addr = await this.start()
+      } catch (e) {
+          if (this.#retryOnFailure && this.#maxNumberOfIntents > this.#intent) {
+              this.#intent++
+              return this.launch()
+          }
+          throw new RunmeServerError(`Cannot start server. Error: ${(e as Error).message}`)
+      }
+
+      await this.connect()
+
+      return addr
+    }
+
+    protected async connect(): Promise<void> {
+      this.closeTransport()
+      await this.acceptsConnection()
+
+      this.#onTransportReady.fire({ transport: this.transport() })
     }
 
     private _port() {
       return this.#port
     }
+
+    protected register<T extends Disposable>(disposable: T): T {
+      this.#disposables.push(disposable)
+      return disposable
+    }
 }
 
 export default RunmeServer


Test patch:
diff --git a/tests/extension/extension.test.ts b/tests/extension/extension.test.ts
index 08a7a94e..63ee6b86 100644
--- a/tests/extension/extension.test.ts
+++ b/tests/extension/extension.test.ts
@@ -6,10 +6,23 @@ import { RunmeExtension } from '../../src/extension/extension'
 vi.mock('vscode')
 vi.mock('vscode-telemetry')
 
-vi.mock('../../src/extension/grpc/client', () => ({
-  ParserServiceClient: vi.fn(),
-  RunnerServiceClient: vi.fn(),
-}))
+vi.mock('../../src/extension/grpc/client', () => {
+  class MockedParserServiceClient {
+    deserialize = vi.fn(() => {
+      return { status: { code: 'OK' } }
+    })
+  }
+
+  return ({
+    ParserServiceClient: MockedParserServiceClient,
+    RunnerServiceClient: vi.fn(),
+    initParserClient: vi.fn(() => ({
+      deserialize: vi.fn(() => {
+        return { status: { code: 'OK' } }
+      })
+    })),
+  })
+})
 
 vi.mock('../../src/extension/grpc/runnerTypes', () => ({}))
 
diff --git a/tests/extension/runner.test.ts b/tests/extension/runner.test.ts
index 8881e65d..819abd87 100644
--- a/tests/extension/runner.test.ts
+++ b/tests/extension/runner.test.ts
@@ -1,7 +1,9 @@
 import path from 'node:path'
 
-import { vi, suite, test, expect } from 'vitest'
+import { vi, suite, test, expect, beforeEach } from 'vitest'
 import type { ExecuteResponse } from '@buf/stateful_runme.community_timostamm-protobuf-ts/runme/runner/v1/runner_pb'
+import { type GrpcTransport } from '@protobuf-ts/grpc-transport'
+import { EventEmitter } from 'vscode'
 
 import {
   GrpcRunner,
@@ -14,9 +16,8 @@ vi.mock('../../src/extension/utils', () => ({
   getGrpcHost: vi.fn().mockReturnValue('127.0.0.1:7863')
 }))
 
-vi.mock('vscode', () => ({
-  ...import(path.join(process.cwd(), '__mocks__', 'vscode')) ,
-  EventEmitter: getEventEmitterClass()
+vi.mock('vscode', async () => ({
+  ...await import(path.join(process.cwd(), '__mocks__', 'vscode')) ,
 }))
 
 vi.mock('@protobuf-ts/grpc-transport', () => ({
@@ -39,29 +40,6 @@ const deleteSession = vi.fn(async () => ({
 
 }))
 
-function getEventEmitterClass() {
-  return class EventEmitter<T> {
-    listeners: MessageCallback<T>[] = []
-
-    event: Event<T> = (listener) => {
-      this.listeners.push(listener)
-
-      return () => {
-        this.listeners = this.listeners.filter(x => x !== listener)
-      }
-    }
-
-    fire(data: T) {
-      this.listeners.forEach(l => l(data))
-    }
-  }
-}
-
-type MessageCallback<T> = (message: T) => void
-type Event<T> = (listener: MessageCallback<T>) => (() => void)
-
-const EventEmitter = getEventEmitterClass()
-
 class MockedDuplexClientStream {
   constructor() {}
 
@@ -112,12 +90,22 @@ vi.mock('@buf/stateful_runme.community_timostamm-protobuf-ts/runme/runner/v1/run
   }
 }))
 
+class MockedRunmeServer {
+  _onTransportReady = new EventEmitter<{ transport: GrpcTransport }>()
+  _onClose = new EventEmitter<{ code: number|null }>()
+
+  onTransportReady = this._onTransportReady.event
+  onClose = this._onClose.event
+}
+
+beforeEach(() => {
+  resetId()
+  deleteSession.mockClear()
+})
+
 suite('grpc Runner', () => {
   test('environment dispose is called on runner dispose', async () => {
-    resetId()
-    deleteSession.mockClear()
-
-    const runner = new GrpcRunner()
+    const { runner } = createGrpcRunner()
     const environment = (await runner.createEnvironment()) as GrpcRunnerEnvironment
 
     const oldEnvDispose = environment.dispose
@@ -134,6 +122,42 @@ suite('grpc Runner', () => {
     expect(deleteSession).toBeCalledTimes(1)
   })
 
+  test('cannot create environment if server not initialized', async () => {
+    const { runner } = createGrpcRunner(false)
+    await expect(runner.createEnvironment()).rejects.toThrowError('Client is not active!')
+  })
+
+  test('cannot create program session if server not initialized', async () => {
+    const { runner } = createGrpcRunner(false)
+    await expect(runner.createProgramSession({ programName: 'sh' })).rejects.toThrowError('Client is not active!')
+  })
+
+  test('cannot get environment variables not initialized', async () => {
+    const { runner } = createGrpcRunner(false)
+    await expect(runner.getEnvironmentVariables({} as any)).rejects.toThrowError('Client is not active!')
+  })
+
+  test('cannot create environment if server closed', async () => {
+    const { runner, server } = createGrpcRunner(true)
+
+    server._onClose.fire({ code: null })
+    await expect(runner.createEnvironment()).rejects.toThrowError('Client is not active!')
+  })
+
+  test('cannot create program session if server closed', async () => {
+    const { runner, server } = createGrpcRunner(true)
+
+    server._onClose.fire({ code: null })
+    await expect(runner.createProgramSession({ programName: 'sh' })).rejects.toThrowError('Client is not active!')
+  })
+
+  test('cannot get environment variables if server closed', async () => {
+    const { runner, server } = createGrpcRunner(true)
+
+    server._onClose.fire({ code: null })
+    await expect(runner.getEnvironmentVariables({} as any)).rejects.toThrowError('Client is not active!')
+  })
+
   suite('grpc program session', () => {
     test('session dispose is called on runner dispose', async () => {
       const { runner, session, duplex } = await createNewSession()
@@ -233,8 +257,21 @@ function getMockedDuplex(session: GrpcRunnerProgramSession): MockedDuplexClientS
   return session['session'] as unknown as MockedDuplexClientStream
 }
 
-async function createNewSession(options: Partial<RunProgramOptions> = {}, runner?: GrpcRunner) {
-  runner ??= new GrpcRunner()
+function createGrpcRunner(initialize = true) {
+  const server = new MockedRunmeServer()
+  const runner = new GrpcRunner(server as any)
+
+  if (initialize) {
+    server._onTransportReady.fire({ transport: {} as any })
+  }
+
+  return { server, runner }
+}
+
+async function createNewSession(options: Partial<RunProgramOptions> = {}, runner?: GrpcRunner, initialize?: boolean) {
+  const { runner: generatedRunner, server } = createGrpcRunner(initialize)
+
+  runner ??= generatedRunner
   const session = (await runner.createProgramSession({
     programName: 'sh',
     ...options
@@ -260,5 +297,6 @@ async function createNewSession(options: Partial<RunProgramOptions> = {}, runner
     closeListener,
     writeListener,
     errListener,
+    server,
   }
 }
diff --git a/tests/extension/server/runmeServer.test.ts b/tests/extension/server/runmeServer.test.ts
index b98180d5..4b9a8e61 100644
--- a/tests/extension/server/runmeServer.test.ts
+++ b/tests/extension/server/runmeServer.test.ts
@@ -43,7 +43,8 @@ suite('Runme server spawn process', () => {
           {
             retryOnFailure: true,
             maxNumberOfIntents: 2,
-          }
+          },
+          false
         )
         const serverLaunchSpy = vi.spyOn(server, 'launch')
         await expect(server.launch()).rejects.toBeInstanceOf(RunmeServerError)
@@ -56,7 +57,8 @@ suite('Runme server spawn process', () => {
         {
           retryOnFailure: true,
           maxNumberOfIntents: 2,
-        }
+        },
+        false
       )
 
       vi.mocked(fs.access).mockResolvedValueOnce()
@@ -84,7 +86,8 @@ suite('Runme server accept connections', () => {
               intents: 4,
               interval: 1,
             }
-          }
+          },
+          false
         )
     })