Refactor connection context

Author: tarrencev

packages/keychain/global.d.ts

@@ -1,4 +1,4 @@
-import { Controller } from "@/utils/controller";
+import Controller from "./src/utils/controller";
 
 declare global {
   interface Window {

packages/keychain/src/components/DeployController.tsx

@@ -21,6 +21,7 @@ import { ControllerError } from "@/utils/connection";
 import { TransactionSummary } from "@/components/transaction/TransactionSummary";
 import { ETH_CONTRACT_ADDRESS, useERC20Balance } from "@cartridge/utils";
 import { Link } from "react-router-dom";
+import { useController } from "@/hooks/controller";
 
 export function DeployController({
   onClose,
@@ -29,7 +30,8 @@ export function DeployController({
   onClose: () => void;
   ctrlError?: ControllerError;
 }) {
-  const { controller, chainName, hasPrefundRequest } = useConnection();
+  const { chainName, hasPrefundRequest } = useConnection();
+  const { controller } = useController();
   const { deploySelf, isDeploying } = useDeploy();
   const [deployHash, setDeployHash] = useState<string>();
   const [error, setError] = useState<Error>();

packages/keychain/src/components/ExecutionContainer.tsx

@@ -1,7 +1,6 @@
 import { useState, useCallback, useEffect, useRef } from "react";
 import { Button } from "@cartridge/ui-next";
 import { Container, Footer } from "@/components/layout";
-import { useConnection } from "@/hooks/connection";
 import { ControllerError } from "@/utils/connection";
 import { ControllerErrorAlert, ErrorAlert } from "@/components/ErrorAlert";
 import { Fees } from "./Fees";
@@ -11,6 +10,7 @@ import { ErrorCode } from "@cartridge/account-wasm/controller";
 import { BannerProps } from "./layout/container/header/Banner";
 import { parseControllerError } from "@/utils/connection/execute";
 import isEqual from "lodash/isEqual";
+import { useController } from "@/hooks/controller";
 
 interface ExecutionContainerProps {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -41,7 +41,7 @@ export function ExecutionContainer({
   children,
 }: ExecutionContainerProps &
   Pick<BannerProps, "title" | "description" | "icon">) {
-  const { controller } = useConnection();
+  const { controller } = useController();
   const [maxFee, setMaxFee] = useState<bigint | null>(null);
   const [ctrlError, setCtrlError] = useState<ControllerError | undefined>(
     executionError,

packages/keychain/src/components/app.tsx

@@ -8,8 +8,25 @@ import { Pending } from "./pending";
 import { Consent, Slot } from "./slot";
 import { OcclusionDetector } from "./OcclusionDetector";
 import { Fund } from "./slot/fund";
+import { useConnection } from "@/hooks/connection";
+import { useController } from "@/hooks/controller";
+import { CreateController } from "./connect";
+import { LoginMode } from "./connect/types";
 
 export function App() {
+  const { origin } = useConnection();
+  const { controller } = useController();
+
+  // Only render in iframe with valid origin
+  if (window.self === window.top || !origin) {
+    return <></>;
+  }
+
+  // If no controller is set up, show the create controller screen
+  if (!controller) {
+    return <CreateController loginMode={LoginMode.Controller} />;
+  }
+
   return (
     <>
       <OcclusionDetector />

packages/keychain/src/components/connect/CreateSession.tsx

@@ -11,6 +11,7 @@ import { UnverifiedSessionSummary } from "@/components/session/UnverifiedSession
 import { VerifiedSessionSummary } from "@/components/session/VerifiedSessionSummary";
 import { DEFAULT_SESSION_DURATION, NOW } from "@/const";
 import { Button, Checkbox } from "@cartridge/ui-next";
+import { useController } from "@/hooks/controller";
 
 export function CreateSession({
   policies,
@@ -21,7 +22,8 @@ export function CreateSession({
   onConnect: (transaction_hash?: string, expiresAt?: bigint) => void;
   isUpdate?: boolean;
 }) {
-  const { controller, upgrade, chainId, theme, logout } = useConnection();
+  const { upgrade, chainId, theme, logout } = useConnection();
+  const { controller } = useController();
   const [isConnecting, setIsConnecting] = useState(false);
   const [isConsent, setIsConsent] = useState(false);
   const [duration, setDuration] = useState<bigint>(DEFAULT_SESSION_DURATION);

packages/keychain/src/components/connect/RegisterSession.tsx

@@ -11,6 +11,7 @@ import { DEFAULT_SESSION_DURATION, NOW } from "@/const";
 import { UnverifiedSessionSummary } from "@/components/session/UnverifiedSessionSummary";
 import { VerifiedSessionSummary } from "@/components/session/VerifiedSessionSummary";
 import { ParsedSessionPolicies } from "@/hooks/session";
+import { useController } from "@/hooks/controller";
 
 export function RegisterSession({
   policies,
@@ -21,7 +22,8 @@ export function RegisterSession({
   onConnect: (transaction_hash: string, expiresAt: bigint) => void;
   publicKey?: string;
 }) {
-  const { controller, theme } = useConnection();
+  const { theme } = useConnection();
+  const { controller } = useController();
   const [duration, setDuration] = useState<bigint>(DEFAULT_SESSION_DURATION);
   const [transactions, setTransactions] = useState<
     | {

packages/keychain/src/components/connect/Upgrade.tsx

@@ -2,9 +2,11 @@ import { BoltIcon, CircleIcon } from "@cartridge/ui-next";
 import { ExecutionContainer } from "@/components/ExecutionContainer";
 import { Content } from "@/components/layout";
 import { useConnection } from "@/hooks/connection";
+import { useController } from "@/hooks/controller";
 
 export const Upgrade = () => {
-  const { upgrade, controller } = useConnection();
+  const { upgrade } = useConnection();
+  const { controller } = useController();
 
   return (
     <ExecutionContainer

packages/keychain/src/components/connect/create/useCreateController.ts

@@ -7,6 +7,7 @@ import Controller from "@/utils/controller";
 import { fetchAccount } from "./utils";
 import { PopupCenter } from "@/utils/url";
 import { useAccountQuery } from "@cartridge/utils/api/cartridge";
+import { useController } from "@/hooks/controller";
 
 export function useCreateController({
   onCreated,
@@ -20,7 +21,8 @@ export function useCreateController({
   const [isLoading, setIsLoading] = useState(false);
   const [error, setError] = useState<Error>();
   const [pendingUsername, setPendingUsername] = useState<string>();
-  const { origin, policies, chainId, rpcUrl, setController } = useConnection();
+  const { origin, policies, chainId, rpcUrl } = useConnection();
+  const { setController } = useController();
 
   useAccountQuery(
     { username: pendingUsername || "" },
@@ -77,7 +79,6 @@ export function useCreateController({
         credentialId,
       });
 
-      window.controller = controller;
       setController(controller);
       onCreated?.();
     },

packages/keychain/src/components/funding/DepositEth.tsx

@@ -36,6 +36,7 @@ import { TokenPair, usePriceQuery } from "@cartridge/utils/api/cartridge";
 import { ETH_CONTRACT_ADDRESS } from "@cartridge/utils";
 import { toast } from "sonner";
 import { DEFAULT_AMOUNT } from "./constants";
+import { useController } from "@/hooks/controller";
 
 type DepositEthProps = {
   onComplete?: (deployHash?: string) => void;
@@ -52,7 +53,8 @@ export function DepositEth(innerProps: DepositEthProps) {
 
 function DepositEthInner({ onComplete, onBack }: DepositEthProps) {
   const { connectAsync, connectors, isPending: isConnecting } = useConnect();
-  const { controller, chainId } = useConnection();
+  const { chainId } = useConnection();
+  const { controller } = useController();
   const { account: extAccount } = useAccount();
 
   const [dollarAmount, setDollarAmount] = useState<number>(DEFAULT_AMOUNT);
@@ -245,7 +247,7 @@ function DepositEthInner({ onComplete, onBack }: DepositEthProps) {
 
 function ExternalWalletProvider({ children }: PropsWithChildren) {
   const { connectors } = useInjectedConnectors({});
-  const { controller } = useConnection();
+  const { controller } = useController();
 
   if (!controller) {
     return children;

packages/keychain/src/components/funding/PurchaseCredits.tsx

@@ -16,6 +16,7 @@ import { Balance } from "./Balance";
 import CheckoutForm from "./StripeCheckout";
 import { isIframe } from "@cartridge/utils";
 import { DEFAULT_AMOUNT } from "./constants";
+import { useController } from "@/hooks/controller";
 
 const STRIPE_API_PUBKEY =
   "pk_test_51Kr6IXIS6lliDpf33KnwWDtIjRPWt3eAI9CuSLR6Vvc3GxHEwmSU0iszYbUlgUadSRluGKAFphe3JzltyjPAKiBK00al4RAFQu";
@@ -32,7 +33,8 @@ type PurchaseCreditsProps = {
 };
 
 export function PurchaseCredits({ isSlot, onBack }: PurchaseCreditsProps) {
-  const { controller, closeModal } = useConnection();
+  const { closeModal } = useConnection();
+  const { controller } = useController();
 
   const [clientSecret, setClientSecret] = useState("");
   const [isLoading, setisLoading] = useState<boolean>(false);

packages/keychain/src/components/funding/index.tsx

@@ -1,6 +1,5 @@
 import { Container, Content, Footer } from "@/components/layout";
 import { useState } from "react";
-import { useConnection } from "@/hooks/connection";
 import {
   Button,
   ArrowIcon,
@@ -11,6 +10,7 @@ import {
 import { DepositEth } from "./DepositEth";
 import { PurchaseCredits } from "./PurchaseCredits";
 import { Balance, BalanceType } from "./Balance";
+import { useController } from "@/hooks/controller";
 
 const enum FundingState {
   SHOW_OPTIONS,
@@ -25,7 +25,7 @@ export type FundingProps = {
 };
 
 export function Funding({ title, isSlot, onComplete }: FundingProps) {
-  const { controller } = useConnection();
+  const { controller } = useController();
   const [state, setState] = useState<FundingState>(FundingState.SHOW_OPTIONS);
   const showBalances: BalanceType[] = isSlot ? ["credits"] : ["credits", "eth"];
   const showCredits =

packages/keychain/src/components/home.tsx

@@ -4,25 +4,23 @@ import { ResponseCodes } from "@cartridge/controller";
 import { useConnection } from "@/hooks/connection";
 import { DeployCtx, ExecuteCtx, SignMessageCtx } from "@/utils/connection";
 import { ConfirmTransaction } from "./transaction/ConfirmTransaction";
-import { CreateController, CreateSession, Logout, Upgrade } from "./connect";
-import { LoginMode } from "./connect/types";
+import { CreateSession, Logout, Upgrade } from "./connect";
 import { DeployController } from "./DeployController";
 import { ErrorPage } from "./ErrorBoundary";
 import { PurchaseCredits } from "./funding/PurchaseCredits";
 import { Settings } from "./settings";
 import { SignMessage } from "./SignMessage";
 import { PageLoading } from "./Loading";
 import { execute } from "@/utils/connection/execute";
-import { usePostHog } from "@/hooks/posthog";
+import { useController } from "@/hooks/controller";
 
 export function Home() {
-  const { context, setContext, controller, error, policies, upgrade } =
-    useConnection();
+  const { context, setContext, error, policies, upgrade } = useConnection();
+  const { controller } = useController();
+
   const [hasSessionForPolicies, setHasSessionForPolicies] = useState<
     boolean | undefined
   >(undefined);
-  const posthog = usePostHog();
-
   useEffect(() => {
     if (controller && policies) {
       controller.isRequestedSession(policies).then((isRequestedSession) => {
@@ -33,29 +31,19 @@ export function Home() {
     }
   }, [controller, policies]);
 
-  useEffect(() => {
-    if (context?.type) {
-      posthog?.capture(
-        `Call ${context.type.charAt(0).toUpperCase() + context.type.slice(1)}`,
-      );
-    }
-  }, [context?.type, posthog]);
-
-  if (window.self === window.top || !context?.origin) {
-    return <></>;
-  }
-
   if (error) {
     return <ErrorPage error={error} />;
   }
 
-  // No controller, send to login
+  if (!context) {
+    return <></>;
+  }
+
   if (!controller) {
-    return <CreateController loginMode={LoginMode.Controller} />;
+    return <PageLoading />;
   }
 
   if (!upgrade.isSynced || hasSessionForPolicies === undefined) {
-    // This is likely never observable in a real application but just in case.
     return <PageLoading />;
   }
 
@@ -113,7 +101,12 @@ export function Home() {
         <SignMessage
           origin={ctx.origin}
           typedData={ctx.typedData}
-          onSign={(sig: Signature) => context.resolve(sig)}
+          onSign={(sig: Signature) =>
+            ctx.resolve({
+              code: ResponseCodes.SUCCESS,
+              ...sig,
+            })
+          }
           onCancel={() =>
             ctx.resolve({
               code: ResponseCodes.CANCELED,

packages/keychain/src/components/provider/connection.tsx

@@ -1,5 +1,4 @@
 import { createContext } from "react";
-import Controller from "@/utils/controller";
 import { ConnectionCtx } from "@/utils/connection";
 import { UpgradeInterface } from "@/hooks/upgrade";
 import { ParsedSessionPolicies } from "@/hooks/session";
@@ -11,7 +10,6 @@ export const ConnectionContext = createContext<
 
 export type ConnectionContextValue = {
   context: ConnectionCtx;
-  controller?: Controller;
   origin?: string;
   rpcUrl?: string;
   chainId?: string;
@@ -22,7 +20,6 @@ export type ConnectionContextValue = {
   error?: Error;
   upgrade: UpgradeInterface;
   setContext: (context: ConnectionCtx) => void;
-  setController: (controller: Controller) => void;
   closeModal: () => void;
   openModal: () => void;
   logout: () => void;

packages/keychain/src/components/provider/posthog.tsx

@@ -1,15 +1,17 @@
 import { PropsWithChildren, useEffect } from "react";
 import { PostHogContext, PostHogWrapper } from "@/context/posthog";
 import { useConnectionValue } from "@/hooks/connection";
+import { useController } from "@/hooks/controller";
 
-export function PostHogProvider({ children }: PropsWithChildren) {
-  const { controller, origin } = useConnectionValue();
+export const posthog = new PostHogWrapper(import.meta.env.VITE_POSTHOG_KEY!, {
+  host: import.meta.env.VITE_POSTHOG_HOST,
+  persistence: "memory",
+  autocapture: false,
+});
 
-  const posthog = new PostHogWrapper(import.meta.env.VITE_POSTHOG_KEY!, {
-    host: import.meta.env.VITE_POSTHOG_HOST,
-    persistence: "memory",
-    autocapture: false,
-  });
+export function PostHogProvider({ children }: PropsWithChildren) {
+  const { origin } = useConnectionValue();
+  const { controller } = useController();
 
   useEffect(() => {
     if (controller) {
@@ -21,13 +23,13 @@ export function PostHogProvider({ children }: PropsWithChildren) {
     } else {
       posthog.reset();
     }
-  }, [posthog, controller]);
+  }, [controller]);
 
   useEffect(() => {
     if (origin) {
       posthog.group("company", origin);
     }
-  }, [posthog, origin]);
+  }, [origin]);
 
   return (
     <PostHogContext.Provider value={{ posthog }}>

packages/keychain/src/components/session.tsx

@@ -11,6 +11,7 @@ import { useCallback, useEffect, useMemo, useState } from "react";
 import { LoginMode } from "@/components/connect/types";
 import { PageLoading } from "@/components/Loading";
 import { useNavigate, useSearchParams } from "react-router-dom";
+import { useController } from "@/hooks/controller";
 
 type SessionResponse = {
   username: string;
@@ -44,7 +45,8 @@ export function Session() {
     [searchParams],
   );
 
-  const { controller, policies, origin } = useConnection();
+  const { policies, origin } = useConnection();
+  const { controller } = useController();
   const [isLoading, setIsLoading] = useState<boolean>(true);
 
   // Handler for calling the callback uri.

packages/keychain/src/components/settings/Delegate.tsx

@@ -4,9 +4,11 @@ import { useConnection } from "@/hooks/connection";
 import { useCallback, useEffect, useState } from "react";
 import { CallData, num } from "starknet";
 import { ExecuteCtx } from "@/utils/connection";
+import { useController } from "@/hooks/controller";
 
 export function Delegate({ onBack }: { onBack: () => void }) {
-  const { controller, context, setContext } = useConnection();
+  const { context, setContext } = useConnection();
+  const { controller } = useController();
   const [delegateAddress, setDelegateAddress] = useState("");
   const [isValid, setIsValid] = useState(true);