type ExtractContextShape<TContext> = TContext extends ContextDefinition<
  infer S,
  Record<string, unknown>
>
  ? S
  : never

const context = defineContext(
  z.object({
    userId: z.string(),
    age: z.number(),
  })
);

type ContextShape = ExtractContextShape<typeof context>;
// ContextShape = { userId: z.ZodString; age: z.ZodNumber }

type ExtractContextTools<TContext> = TContext extends ContextDefinition<
  z.ZodRawShape,
  infer TTools
>
  ? TTools
  : Record<string, unknown>

const context = defineContext(
  z.object({ userId: z.string() }),
  {
    tools: {
      logger: console,
      database: dbClient,
    },
  }
);

type ContextTools = ExtractContextTools<typeof context>;
// ContextTools = { logger: Console; database: DatabaseClient }