From 45eb5b437a690b6a0d33fea3174d53ed02f88380 Mon Sep 17 00:00:00 2001
From: blam <ben@blam.sh>
Date: Thu, 1 Aug 2024 09:41:40 +0200
Subject: [PATCH] chore: got the data container working nicely

Signed-off-by: blam <ben@blam.sh>
---
 .../src/wiring/createExtension.test.ts        |  3 ++
 .../src/wiring/createExtension.ts             | 28 +++----------------
 2 files changed, 7 insertions(+), 24 deletions(-)

diff --git a/packages/frontend-plugin-api/src/wiring/createExtension.test.ts b/packages/frontend-plugin-api/src/wiring/createExtension.test.ts
index e38a24cebe..c688f92004 100644
--- a/packages/frontend-plugin-api/src/wiring/createExtension.test.ts
+++ b/packages/frontend-plugin-api/src/wiring/createExtension.test.ts
@@ -665,6 +665,9 @@ describe('createExtension', () => {
 
           const foo: string = response.get(stringDataRef);
 
+          // @ts-expect-error - fails because original factory does not return number
+          const number: boolean = response.get(numberDataRef);
+
           return [stringDataRef(`foo-${foo}-override`)];
         },
       });
diff --git a/packages/frontend-plugin-api/src/wiring/createExtension.ts b/packages/frontend-plugin-api/src/wiring/createExtension.ts
index 678ce4e042..87a4342df7 100644
--- a/packages/frontend-plugin-api/src/wiring/createExtension.ts
+++ b/packages/frontend-plugin-api/src/wiring/createExtension.ts
@@ -20,7 +20,6 @@ import { Expand } from '../types';
 import {
   AnyExtensionDataRef,
   ExtensionDataRef,
-  ExtensionDataRefToValue,
   ExtensionDataValue,
 } from './createExtensionDataRef';
 import { ExtensionInput, LegacyExtensionInput } from './createExtensionInput';
@@ -246,7 +245,6 @@ export type OverrideExtensionOptions<
   TConfig,
   TConfigInput,
   TConfigSchema extends { [key: string]: (zImpl: typeof z) => z.ZodType },
-  UOriginalFactoryOutput extends ExtensionDataValue<any, any>,
   TConfigSchemaOverrides extends {
     [key: string]: (zImpl: typeof z) => z.ZodType;
   },
@@ -274,9 +272,7 @@ export type OverrideExtensionOptions<
         [key in keyof TConfigSchema]: z.infer<ReturnType<TConfigSchema[key]>>;
       };
       inputs?: Expand<ResolvedExtensionInputs<TInputs>>;
-      // todo(blam): Think this is better as a DataContainer instead
-      // should probably update that everywhere.
-    }) => Iterable<UOriginalFactoryOutput>,
+    }) => ExtensionDataContainer<UOutput>,
     context: {
       node: AppNode;
       config: TConfig &
@@ -296,7 +292,7 @@ export type OverrideExtensionOptions<
   ): Iterable<UFactoryOverrideOutput>;
   // todo(blam): need to verify that the outputs are merged and verified properly.
 } & VerifyExtensionFactoryOutput<
-  UOutput & UFactoryOverrideOutput,
+  UOutput & UOutputOverrides,
   UFactoryOverrideOutput
 >;
 
@@ -312,7 +308,6 @@ export type OverridableExtension<
   TConfig,
   TConfigInput,
   TConfigSchema extends { [key: string]: (zImpl: typeof z) => z.ZodType },
-  UFactoryOutput extends ExtensionDataValue<any, any>,
 > = {
   override<
     TConfigSchemaOverrides extends {
@@ -327,7 +322,6 @@ export type OverridableExtension<
       TConfig,
       TConfigInput,
       TConfigSchema,
-      UFactoryOutput,
       TConfigSchemaOverrides,
       UOutputOverrides,
       UFactoryOverrideOutput
@@ -431,14 +425,7 @@ export function createExtension<
           }>
         >)
 > &
-  OverridableExtension<
-    UOutput,
-    TInputs,
-    TConfig,
-    TConfigInput,
-    TConfigSchema,
-    UFactoryOutput
-  >;
+  OverridableExtension<UOutput, TInputs, TConfig, TConfigInput, TConfigSchema>;
 /**
  * @public
  * @deprecated - use the array format of `output` instead, see TODO-doc-link
@@ -519,14 +506,7 @@ export function createExtension<
           }>
         >)
 > &
-  OverridableExtension<
-    UOutput,
-    TInputs,
-    TConfig,
-    TConfigInput,
-    TConfigSchema,
-    UFactoryOutput
-  > {
+  OverridableExtension<UOutput, TInputs, TConfig, TConfigInput, TConfigSchema> {
   const newConfigSchema = options.config?.schema;
   if (newConfigSchema && options.configSchema) {
     throw new Error(`Cannot provide both configSchema and config.schema`);