Merge pull request #276 from smartcontractkit/rtinianov_teeRuntime

Add the TEERuntime

Author: nolag

.gitmodules

@@ -1,4 +1,4 @@
 [submodule "submodules/chainlink-protos"]
 	path = submodules/chainlink-protos
 	url = https://github.com/smartcontractkit/chainlink-protos.git
-	branch = capabilities-development
+	branch = capabilities-development

packages/cre-sdk/src/generated-sdk/capabilities/networking/http/v1alpha/client_sdk_gen.ts

@@ -6,7 +6,7 @@ import {
 	type Response,
 	ResponseSchema,
 } from '@cre/generated/capabilities/networking/http/v1alpha/client_pb'
-import type { NodeRuntime, Runtime } from '@cre/sdk'
+import type { NodeRuntime, Runtime, TeeRuntime } from '@cre/sdk'
 import { Report } from '@cre/sdk/report'
 import type { ConsensusAggregation, PrimitiveTypes, UnwrapOptions } from '@cre/sdk/utils'
 import type { CapabilityInput } from '@cre/sdk/utils/types/no-excess'
@@ -37,10 +37,18 @@ export class ClientCapability {
 	static readonly CAPABILITY_NAME = 'http-actions'
 	static readonly CAPABILITY_VERSION = '1.0.0-alpha'
 
+	sendRequest<TInput>(
+		runtime: TeeRuntime<unknown>,
+		input: CapabilityInput<TInput, Request, RequestJson>,
+	): { result: () => Response }
 	sendRequest<TInput>(
 		runtime: NodeRuntime<unknown>,
 		input: CapabilityInput<TInput, Request, RequestJson>,
 	): { result: () => Response }
+	sendRequest<TInput>(
+		runtime: NodeRuntime<unknown> | TeeRuntime<unknown>,
+		input: CapabilityInput<TInput, Request, RequestJson>,
+	): { result: () => Response }
 	sendRequest<TArgs extends unknown[], TOutput>(
 		runtime: Runtime<unknown>,
 		fn: (sendRequester: SendRequester, ...args: TArgs) => TOutput,
@@ -59,11 +67,14 @@ export class ClientCapability {
 			return this.sendRequestSugarHelper(runtime, fn, consensusAggregation, unwrapOptions)
 		}
 		// Otherwise, this is the basic call overload
-		const [runtime, input] = args as [NodeRuntime<unknown>, Request | RequestJson]
+		const [runtime, input] = args as [
+			NodeRuntime<unknown> | TeeRuntime<unknown>,
+			Request | RequestJson,
+		]
 		return this.sendRequestCallHelper(runtime, input)
 	}
 	private sendRequestCallHelper(
-		runtime: NodeRuntime<unknown>,
+		runtime: NodeRuntime<unknown> | TeeRuntime<unknown>,
 		input: Request | RequestJson,
 	): { result: () => Response } {
 		// Handle input conversion - unwrap if it's a wrapped type, convert from JSON if needed

packages/cre-sdk/src/generator/generate-action.ts

@@ -17,6 +17,7 @@ export function generateActionMethod(
 	capabilityClassName: string,
 	labels: ProcessedLabel[],
 	modePrefix: string,
+	teeEnabled: boolean,
 ): string {
 	const capabilityIdLogic = generateCapabilityIdLogic(labels, capabilityClassName)
 
@@ -47,9 +48,25 @@ export function generateActionMethod(
 	//   - JSON shape -> wrapped in NoExcess so unknown keys (e.g. plain
 	//     `body` instead of `bodyString`) fail at the call boundary even
 	//     when the user lifts the request object into a variable.
-	const callSig = `<TInput>(runtime: ${modePrefix}Runtime<unknown>, input: CapabilityInput<TInput, ${nativeInputType}, ${jsonInputType}>): {result: () => ${outputType}}`
-	// Internal impl signature - widest, accepts either form.
-	const implSig = `(runtime: ${modePrefix}Runtime<unknown>, input: ${nativeInputType} | ${jsonInputType}): {result: () => ${outputType}}`
+	const basicSig = `<TInput>(runtime: ${modePrefix}Runtime<unknown>, input: CapabilityInput<TInput, ${nativeInputType}, ${jsonInputType}>): {result: () => ${outputType}}`
+
+	const callSig = teeEnabled
+		? basicSig.replace(
+				`${modePrefix}Runtime<unknown>`,
+				`${modePrefix}Runtime<unknown> | TeeRuntime<unknown>`,
+			)
+		: basicSig
+
+	const teeSig = basicSig.replace(`${modePrefix}Runtime<unknown>`, `TeeRuntime<unknown>`)
+
+	var nameAndPublicSigs = teeEnabled
+		? `${methodName}${teeSig}\n${methodName}${basicSig};\n${methodName}${callSig};`
+		: `${methodName}${callSig};`
+
+	// Internal impl signature - widest, accepts either form (and either runtime when TEE is enabled).
+	const implSig = `(runtime: ${modePrefix}Runtime<unknown>${
+		teeEnabled ? ' | TeeRuntime<unknown>' : ''
+	}, input: ${nativeInputType} | ${jsonInputType}): {result: () => ${outputType}}`
 	const callSigAndBody = `${implSig} {
     // Handle input conversion - unwrap if it's a wrapped type, convert from JSON if needed
     let payload: ${method.input.name}
@@ -113,7 +130,7 @@ export function generateActionMethod(
       : UnwrapOptions<TOutput>,
   ): (...args: TArgs) => { result: () => TOutput }`
 		return `
-  ${methodName}${callSig}
+  ${nameAndPublicSigs}
   ${methodName}${sugarSig}
   ${methodName}(...args: unknown[]): unknown {
     // Check if this is the sugar syntax overload (has function parameter)
@@ -122,7 +139,7 @@ export function generateActionMethod(
       return this.${methodName}SugarHelper(runtime, fn, consensusAggregation, unwrapOptions)
     }
     // Otherwise, this is the basic call overload
-    const [runtime, input] = args as [${modePrefix}Runtime<unknown>, ${inputTypes.join(' | ')}]
+    const [runtime, input] = args as [${teeEnabled ? `${modePrefix}Runtime<unknown> | TeeRuntime<unknown>` : `${modePrefix}Runtime<unknown>`}, ${inputTypes.join(' | ')}]
     return this.${methodName}CallHelper(runtime, input)
   }
   private ${methodName}CallHelper${callSigAndBody}
@@ -134,7 +151,13 @@ export function generateActionMethod(
       return runtime.runInNodeMode(wrappedFn, consensusAggregation, unwrapOptions)
     }`
 	}
+
+	// For DON mode: emit tee + basic overload declarations, then the implementation with its name.
+	// nameAndPublicSigs is designed for Node mode's dispatcher pattern and must not be reused here.
+	const donOverloads = teeEnabled
+		? `${methodName}${teeSig}\n  ${methodName}${basicSig}\n  `
+		: `${methodName}${callSig}\n  `
+
 	return `
-  ${methodName}${callSig}
-  ${methodName}${callSigAndBody}`
+  ${donOverloads}${methodName}${callSigAndBody}`
 }

packages/cre-sdk/src/generator/generate-sdk.ts

@@ -5,19 +5,15 @@ import type { GenFile } from '@bufbuild/protobuf/codegenv2'
 import { Mode } from '@cre/generated/sdk/v1alpha/sdk_pb'
 import type { CapabilityMetadata } from '@cre/generated/tools/generator/v1alpha/cre_metadata_pb'
 import {
+	AdditionalEnvironments,
 	capability,
 	method as methodOption,
 } from '@cre/generated/tools/generator/v1alpha/cre_metadata_pb'
 import { generateActionMethod } from './generate-action'
 import { generateReportWrapper } from './generate-report-wrapper'
 import { generateActionSugarClass } from './generate-sugar'
 import { generateTriggerClass, generateTriggerMethod } from './generate-trigger'
-import {
-	generateCapabilityIdLogic,
-	generateConstructorParams,
-	generateLabelSupport,
-	processLabels,
-} from './label-utils'
+import { generateConstructorParams, generateLabelSupport, processLabels } from './label-utils'
 import { getImportPathForFile, lowerCaseFirstLetter } from './utils'
 
 const getCapabilityServiceOptions = (service: DescService): CapabilityMetadata | false => {
@@ -96,6 +92,7 @@ export function generateSdk(file: GenFile, outputDir: string) {
 		})
 
 		const modePrefix = capOption.mode === Mode.NODE ? 'Node' : ''
+		const teeEnabled = capOption.additionalEnvironments.includes(AdditionalEnvironments.TEE)
 
 		// Build import statements
 		// Note: protobuf imports are deferred until after report wrappers are processed,
@@ -122,10 +119,12 @@ export function generateSdk(file: GenFile, outputDir: string) {
 
 		if (hasActions) {
 			if (modePrefix !== '') {
-				imports.add(`import type { Runtime, ${modePrefix}Runtime } from "@cre/sdk"`)
+				imports.add(
+					`import type { Runtime, ${modePrefix}Runtime${teeEnabled ? ', TeeRuntime' : ''} } from "@cre/sdk"`,
+				)
 				imports.add(`import { Report } from "@cre/sdk/report"`)
 			} else {
-				imports.add(`import type { Runtime } from "@cre/sdk"`)
+				imports.add(`import type { Runtime ${teeEnabled ? ', TeeRuntime' : ''} } from "@cre/sdk"`)
 				imports.add(`import { Report } from "@cre/sdk/report"`)
 				imports.add(`import { hexToBytes } from "@cre/sdk/utils/hex-utils";`)
 			}
@@ -234,7 +233,14 @@ export function generateSdk(file: GenFile, outputDir: string) {
 				}
 
 				// Generate action method
-				return generateActionMethod(method, methodName, capabilityClassName, labels, modePrefix)
+				return generateActionMethod(
+					method,
+					methodName,
+					capabilityClassName,
+					labels,
+					modePrefix,
+					teeEnabled,
+				)
 			})
 			.join('\n')
 

packages/cre-sdk/src/sdk/cre/index.ts

@@ -10,7 +10,7 @@ import { ClientCapability as HTTPClient } from '@cre/generated-sdk/capabilities/
 import { HTTPCapability } from '@cre/generated-sdk/capabilities/networking/http/v1alpha/http_sdk_gen'
 import { CronCapability } from '@cre/generated-sdk/capabilities/scheduler/cron/v1/cron_sdk_gen'
 import { prepareRuntime } from '@cre/sdk/utils/prepare-runtime'
-import { handler } from '@cre/sdk/workflow'
+import { handler, handlerInTee } from '@cre/sdk/workflow'
 
 /**
  * Public exports for the CRE SDK.
@@ -26,6 +26,7 @@ export {
 } from '@cre/generated/capabilities/blockchain/solana/v1alpha/client_pb'
 export type { Payload as HTTPPayload } from '@cre/generated/capabilities/networking/http/v1alpha/trigger_pb'
 export type { Payload as CronPayload } from '@cre/generated/capabilities/scheduler/cron/v1/trigger_pb'
+export { TeeType } from '@cre/generated/sdk/v1alpha/sdk_pb'
 // Aptos Capability
 export {
 	ClientCapability as AptosClient,
@@ -53,13 +54,20 @@ export {
 	type SendRequester as HTTPSendRequester,
 } from '@cre/generated-sdk/capabilities/networking/http/v1alpha/client_sdk_gen'
 export { HTTPCapability } from '@cre/generated-sdk/capabilities/networking/http/v1alpha/http_sdk_gen'
-
 // CRON Capability
 export { CronCapability } from '@cre/generated-sdk/capabilities/scheduler/cron/v1/cron_sdk_gen'
-
 // Runtime
-export type { NodeRuntime, Runtime } from '@cre/sdk/runtime'
-export { handler } from '@cre/sdk/workflow'
+export type { NodeRuntime, Runtime, TeeRuntime } from '@cre/sdk/runtime'
+export type {
+	AnyTeeConstraint,
+	NitroBinding,
+	NitroRegion,
+	OneOfTees,
+	Region,
+	TeeBinding,
+	TeeConstraint,
+} from '@cre/sdk/workflow'
+export { handler, handlerInTee, NITRO_REGIONS, REGIONS } from '@cre/sdk/workflow'
 
 prepareRuntime()
 
@@ -74,4 +82,5 @@ export const cre = {
 		SolanaClient,
 	},
 	handler,
+	handlerInTee,
 }

packages/cre-sdk/src/sdk/impl/runtime-impl.ts

@@ -34,6 +34,7 @@ import type {
 	ReportRequest,
 	ReportRequestJson,
 	Runtime,
+	TeeRuntime,
 } from '@cre/sdk'
 import type { Report } from '@cre/sdk/report'
 import {
@@ -462,6 +463,63 @@ export class RuntimeImpl<C> extends BaseRuntimeImpl<C> implements Runtime<C> {
 	}
 }
 
+export class TeeRuntimeImpl<C> implements TeeRuntime<C> {
+	private readonly runtime: RuntimeImpl<C>
+	constructor(
+		public config: C,
+		public nextCallId: number,
+		helpers: RuntimeHelpers,
+		maxResponseSize: bigint,
+	) {
+		this.runtime = new RuntimeImpl(config, nextCallId, helpers, maxResponseSize)
+	}
+	getSecret(request: SecretRequest | SecretRequestJson): { result: () => Secret } {
+		return this.runtime.getSecret(request)
+	}
+	now(): Date {
+		return this.runtime.now()
+	}
+
+	log(message: string): void {
+		this.runtime.log(message)
+	}
+
+	emitMetric(
+		name: string,
+		value: number,
+		type: MetricType,
+		labels?: Record<string, string>,
+	): boolean {
+		return this.runtime.emitMetric(name, value, type, labels)
+	}
+
+	callCapability<I extends Message, O extends Message>({
+		capabilityId,
+		method,
+		payload,
+		inputSchema,
+		outputSchema,
+	}: CallCapabilityParams<I, O>): { result: () => O } {
+		return this.runtime.callCapability({
+			capabilityId,
+			method,
+			payload,
+			inputSchema,
+			outputSchema,
+		})
+	}
+
+	reportFromDon(input: ReportRequest | ReportRequestJson): {
+		result: () => Report
+	} {
+		return this.runtime.report(input)
+	}
+
+	usingTheDons(): Runtime<C> {
+		return this.runtime
+	}
+}
+
 /**
  * Interface to the WASM host environment.
  * Provides low-level access to capabilities, secrets, and utilities.

packages/cre-sdk/src/sdk/runtime.ts

@@ -47,6 +47,28 @@ export interface NodeRuntime<C> extends BaseRuntime<C> {
 	readonly _isNodeRuntime: true
 }
 
+/**
+ * Runtime for Tee mode execution.
+ */
+export interface TeeRuntime<C> extends BaseRuntime<C>, SecretsProvider {
+	/**
+	 * Generates a report from the DON.
+	 * Data requestsed throught this method will be routed outside of the TEE.
+	 *
+	 * @param input - Report request to generate a report from the DON
+	 * @returns Report generated from the DON
+	 */
+	reportFromDon(input: ReportRequest | ReportRequestJson): {
+		result: () => Report
+	}
+
+	/**
+	 * Returns the runtime that makes requests to the CRE DONs.
+	 * Requests made through this runtime will therefore be routed outside of the TEE
+	 */
+	usingTheDons(): Runtime<C>
+}
+
 /**
  * Runtime for DON mode execution.
  */