diff --git a/javascript/src/rpc.spec.ts b/javascript/src/rpc.spec.ts
index 943bff0..b5fd6c9 100644
--- a/javascript/src/rpc.spec.ts
+++ b/javascript/src/rpc.spec.ts
@@ -2,11 +2,11 @@ import fs from 'fs-extra'
 import nock from 'nock'
 import crypto from 'crypto'
 import { dedent } from 'ts-dedent'
-import { describe, expect, beforeEach, afterEach, test, vi } from 'vitest'
+import { describe, expect, beforeEach, afterEach, test, vi, expectTypeOf } from 'vitest'
 import { v4 as uuidv4 } from 'uuid'
 
 import { RetoolRPC } from './rpc'
-import { Arguments } from './types'
+import { Arguments, RetoolContext, TransformedArguments } from './types'
 import { parseFunctionArguments } from './utils/schema'
 import { RetoolRPCVersion } from './version'
 
@@ -751,6 +751,25 @@ describe('RetoolRPC', () => {
       )
     })
   })
+
+  test('returns the implementation when registering', async () => {
+    const fn = rpcAgent.register({
+      name: 'test',
+      arguments: {},
+      implementation: async () => {
+        return 1
+      }
+    })
+
+    type ExpectedImplementation = (args: TransformedArguments<Arguments>, context: RetoolContext) => Promise<number>
+    expectTypeOf(fn).toEqualTypeOf<ExpectedImplementation>()
+    
+
+    const result = await fn({}, context);
+    expect(result).toEqual(1)
+    expectTypeOf(result).toEqualTypeOf(1)
+  })
+
 })
 
 describe('RetoolRPCVersion', () => {
diff --git a/javascript/src/rpc.ts b/javascript/src/rpc.ts
index 8470bce..94f1757 100644
--- a/javascript/src/rpc.ts
+++ b/javascript/src/rpc.ts
@@ -36,7 +36,7 @@ export class RetoolRPC {
   private _version: string
   private _agentUuid: string
   private _versionHash: string | undefined
-  private _functions: Record<string, Omit<RegisterFunctionSpec<any>, 'name'>> = {}
+  private _functions: Record<string, Omit<RegisterFunctionSpec<any, any>, 'name'>> = {}
   private _retoolApi: RetoolAPI
   private _logger: LoggerService
 
@@ -89,12 +89,13 @@ export class RetoolRPC {
   /**
    * Registers a Retool function with the specified function definition.
    */
-  register<TArgs extends Arguments>(spec: RegisterFunctionSpec<TArgs>): void {
+  register<TArgs extends Arguments, TReturn>(spec: RegisterFunctionSpec<TArgs, TReturn>): RegisterFunctionSpec<TArgs, TReturn>['implementation'] {
     this._functions[spec.name] = {
       arguments: spec.arguments,
       permissions: spec.permissions,
       implementation: spec.implementation,
     }
+    return spec.implementation;
   }
 
   /**
@@ -145,7 +146,7 @@ export class RetoolRPC {
    * Registers the agent with the Retool server.
    */
   private async registerAgent(): Promise<AgentServerStatus> {
-    const functionsMetadata: Record<string, Pick<RegisterFunctionSpec<any>, 'arguments' | 'permissions'>> = {}
+    const functionsMetadata: Record<string, Pick<RegisterFunctionSpec<any, any>, 'arguments' | 'permissions'>> = {}
     for (const functionName in this._functions) {
       functionsMetadata[functionName] = {
         arguments: this._functions[functionName].arguments,
diff --git a/javascript/src/types.ts b/javascript/src/types.ts
index 8bd66ca..93e3dfa 100644
--- a/javascript/src/types.ts
+++ b/javascript/src/types.ts
@@ -74,13 +74,13 @@ export type TransformedArguments<TArgs extends Arguments> = {
 }
 
 /** Represents the specification for registering a Retool function. */
-export type RegisterFunctionSpec<TArgs extends Arguments> = {
+export type RegisterFunctionSpec<TArgs extends Arguments, TReturn> = {
   /** The name of the function. */
   name: string
   /** The arguments of the function. */
   arguments: Pick<TArgs, keyof TArgs>
   /** The implementation of the function. */
-  implementation: (args: TransformedArguments<TArgs>, context: RetoolContext) => Promise<any>
+  implementation: (args: TransformedArguments<TArgs>, context: RetoolContext) => Promise<TReturn>
   /** The permissions configuration for the function. */
   permissions?: {
     /** The list of group names that have permission to execute the function. */