There are at least two directions in which you can proceed.
First, Church's simple type theory (see also this explanation) extends the simply typed $\lambda$-calculus with the types of natural numbers and booleans, and other constructs that make it sufficiently expressive to encode a certain kind of classical set theory. These ideas were pursued in the HOL theorem provers. For instance, HOL light for a small implementation. This approach has the flavor of doing logic, i.e., there is no general notion of computation - as opposed to the simply typed $\lambda$-calculus, which is a programming language.
Second, dependent type theory, of which there are many variants, is also an extension of the simply typed $\lambda$-calculus. It can be used as a foundation of mathematics. They are used in proof assistants, such as Coq, Agda, and Lean. Some variants, notably Martin-Löf type theory, do retain the character of a programming language. In fact, Idris is a programming language based on dependent type theory.
None of the above is "set theory", but these formalisms are quite suitable for expressing mathematics, including category theory.