Sign in with ChatGPT

npm add @rivet-dev/pi @earendil-works/pi-durable @earendil-works/pi-ai @earendil-works/pi-coding-agent rivetkit

export CHATGPT_HOST_ID=$(node -e 'console.log(crypto.randomUUID())')

import { type Credential, InMemoryCredentialStore } from "@earendil-works/pi-ai";
import { ModelRuntime } from "@earendil-works/pi-coding-agent";
import type { PiProviderCredential } from "@rivet-dev/pi";
import { actor } from "rivetkit";

// One per user. Holds the user's ChatGPT login and gives the agent fresh access tokens.
export const credentials = actor({
	state: { saved: {} as Record<string, Credential> },
	actions: {
		save: (c, provider: string, credential: Credential) => {
			c.state.saved[provider] = credential;
		},
		list: (c) => Object.entries(c.state.saved).map(([providerId, { type }]) => ({ providerId, type })),
		read: (c, provider: string) => withoutRefreshToken(c.state.saved[provider]),
		refresh: async (c, provider: string) => {
			const store = new InMemoryCredentialStore();
			await store.modify(provider, async () => c.state.saved[provider]);
			const runtime = await ModelRuntime.create({ credentials: store, modelsPath: null });
			await runtime.getAuth(provider, { minOAuthValidityMs: 10 * 60_000 });
			const refreshed = await store.read(provider);
			if (refreshed) c.state.saved[provider] = refreshed;
			return withoutRefreshToken(refreshed);
		},
	},
});

function withoutRefreshToken(credential: Credential | undefined): PiProviderCredential | undefined {
	if (credential?.type !== "oauth") return credential;
	const { refresh: _refresh, ...rest } = credential;
	return rest;
}

import { type Credential, createModels, InMemoryCredentialStore } from "@earendil-works/pi-ai";
import { openaiProvider } from "@earendil-works/pi-ai/providers/openai";
import { actor, type Registry, UserError } from "rivetkit";
import type { credentials } from "./credentials";

// A UUID you generate once for this deployment. OpenAI calls it the agent host ID.
const HOST_ID = process.env.CHATGPT_HOST_ID ?? "";
// Your app's name. Users see it on the ChatGPT consent screen.
const APP_NAME = "Acme";

interface PendingSignIn {
	cancel: () => void;
	pasteRedirect: (url: string) => void;
	credential: Promise<Credential>;
}

// One per user. Runs one sign-in at a time and saves the result to the user's credentials Actor.
export const chatgptLogin = actor({
	createVars: () => ({ pending: undefined as PendingSignIn | undefined }),
	actions: {
		// Starts a sign-in and returns the URL to open in the user's browser.
		start: async (c) => {
			c.vars.pending?.cancel();

			const models = createModels({ credentials: new InMemoryCredentialStore() });
			models.setProvider(openaiProvider());

			const cancel = new AbortController();
			const authUrl = Promise.withResolvers<string>();
			const redirect = Promise.withResolvers<string>();
			const credential = models.login(
				"openai",
				"oauth",
				{
					// Give up after ten minutes, or when the user starts over.
					signal: AbortSignal.any([cancel.signal, AbortSignal.timeout(10 * 60_000)]),
					notify: (event) => {
						if (event.type === "auth_url") authUrl.resolve(event.url);
					},
					// pi-ai asks for the redirect URL the browser landed on. Rejecting on abort
					// lets pi-ai close its callback server, so the next sign-in can start.
					prompt: ({ signal }) =>
						new Promise<string>((resolve, reject) => {
							signal?.addEventListener("abort", () => reject(new Error("Sign-in cancelled.")), { once: true });
							redirect.promise.then(resolve);
						}),
				},
				{ getDeviceId: () => HOST_ID, agentName: APP_NAME },
			);

			c.vars.pending = { cancel: () => cancel.abort(), pasteRedirect: redirect.resolve, credential };
			// Stay awake until the sign-in finishes, fails, or times out.
			c.keepAwake(credential.catch(() => undefined));

			// Rejects here if the sign-in fails before it has a URL.
			return Promise.race([authUrl.promise, credential.then(() => authUrl.promise)]);
		},

		// Finishes the sign-in with the redirect URL the user pasted. Without one, it waits
		// for the browser to reach the callback on its own, which works only when the worker
		// runs on the same machine as the browser.
		finish: async (c, redirectUrl?: string) => {
			const pending = c.vars.pending;
			if (!pending) throw new UserError("Start a sign-in first.");
			if (redirectUrl) pending.pasteRedirect(redirectUrl);

			try {
				const credential = await pending.credential;
				const client = c.client<Registry<{ credentials: typeof credentials }>>();
				await client.credentials.getOrCreate([c.key[0]]).save("openai", credential);
			} catch (error) {
				throw new UserError(error instanceof Error ? error.message : "ChatGPT sign-in failed.");
			} finally {
				c.vars.pending = undefined;
			}
		},
	},
});

import { createRegistry } from "@earendil-works/pi-durable";
import { pi } from "@rivet-dev/pi";
import { type Registry, setup } from "rivetkit";
import { chatgptLogin } from "./chatgpt-login";
import { credentials } from "./credentials";

const agent = pi({
	model: "openai/gpt-6-sol",
	registry: createRegistry(),
	// The agent's key starts with the user id, so it runs on that user's ChatGPT plan.
	credentials: (c) => {
		const client = c.client<Registry<{ credentials: typeof credentials }>>();
		return client.credentials.getOrCreate([c.key[0]]);
	},
});

export const registry = setup({ use: { credentials, chatgptLogin, agent } });

registry.start();

import { createInterface } from "node:readline/promises";
import { createClient } from "rivetkit/client";
import type { registry } from "./server";

const client = createClient<typeof registry>();
const userId = "user-123";

// 1. Start the sign-in and send the user to ChatGPT.
const login = client.chatgptLogin.getOrCreate([userId]);
const url = await login.start();
console.log(`Open this URL and choose Continue:\n${url}\n`);

// 2. The browser lands on a 127.0.0.1 URL. The user copies it from the address bar.
const terminal = createInterface({ input: process.stdin, output: process.stdout });
const redirect = await terminal.question("Paste the URL you landed on, or press Enter if the page says you can close it: ");
terminal.close();
await login.finish(redirect.trim() || undefined);

// 3. The agent now runs on the user's ChatGPT plan.
const result = await client.agent.getOrCreate([userId, "chat"]).prompt("Say hi in five words.");
console.log(result.status === "done" ? result.text : `Unanswered: ${result.reason}`);