1import * as GCP from "@/GCP";
2import * as Effect from "effect/Effect";
3import { serveProbes } from "../../bindingHost.ts";
4
5// The gates are decided at deploy time and passed to the deployed host
6// through its env, so the host serves exactly the probes it was bound for.
7const gate = (name: string, atDeploy: boolean) => {
8 const passed = process.env[name];
9 // Host env values arrive JSON-encoded (`"1"`).
10 return passed === undefined ? atDeploy : passed.replace(/"/g, "") === "1";
11};
12
13/**
14 * Sandbox environments and templates take 1-3 minutes to provision.
15 */
16export const runSandboxes = gate("AIPLATFORM_SANDBOXES", !process.env.FAST);
17
18/**
19 * pauseReasoningEnginesSandboxEnvironments does not work on the testing
20 * project (the sandbox itself creates fine): it has answered
21 * InternalServerError "Internal error encountered.", and from the host the
22 * first calls drop the connection (Cloud Run 503) after which the sandbox
23 * answers FAILED_PRECONDITION "is not in RUNNING, PAUSING or STOPPING
24 * state, cannot be paused.". Set GCP_TEST_AIPLATFORM_SANDBOX_PAUSE=1 where
25 * pausing works.
26 */
27export const runSandboxPause = gate(
28 "AIPLATFORM_SANDBOX_PAUSE",
29 runSandboxes && !!process.env.GCP_TEST_AIPLATFORM_SANDBOX_PAUSE,
30);
31
32/** Reasoning engine without a deployed package (queries are refused). */
33export const Agent = GCP.AIPlatform.ReasoningEngine("Agent", {
34 location: "us-central1",
35 displayName: "alchemy-binding-engine",
36 spec: { agentFramework: "custom" },
37});
38
39/** Custom training pipeline the host reads and cancels. */
40export const Train = GCP.AIPlatform.TrainingPipeline("Train", {
41 location: "us-central1",
42 displayName: "alchemy-binding-pipeline",
43 trainingTaskDefinition:
44 "gs://google-cloud-aiplatform/schema/trainingjob/definition/custom_task_1.0.0.yaml",
45 trainingTaskInputs: {
46 workerPoolSpecs: [
47 {
48 machineSpec: { machineType: "n1-standard-4" },
49 replicaCount: "1",
50 containerSpec: {
51 imageUri:
52 "us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-12.py310:latest",
53 // Runs long enough for the host to cancel it.
54 command: ["sleep", "1800"],
55 },
56 },
57 ],
58 },
59});
60
61const sandbox = (id: string, displayName: string) =>
62 Effect.gen(function* () {
63 const engine = yield* Agent;
64 return yield* GCP.AIPlatform.ReasoningEnginesSandboxEnvironment(id, {
65 reasoningEngine: engine.name,
66 displayName,
67 ttl: "600s",
68 spec: {
69 codeExecutionEnvironment: {
70 codeLanguage: "LANGUAGE_PYTHON",
71 machineConfig: "MACHINE_CONFIG_VCPU4_RAM4GIB",
72 },
73 },
74 });
75 });
76
77/** Sandbox the host reads. */
78export const Code = sandbox("Code", "code");
79/** Sandbox the host pauses. */
80export const Paused = sandbox("Paused", "paused");
81/** Sandbox the host pauses and resumes. */
82export const Resumed = sandbox("Resumed", "resumed");
83
84/** Sandbox template the host reads. */
85export const Browser = Effect.gen(function* () {
86 const engine = yield* Agent;
87 return yield* GCP.AIPlatform.ReasoningEnginesSandboxEnvironmentTemplate(
88 "Browser",
89 {
90 reasoningEngine: engine.name,
91 displayName: "browser",
92 defaultContainerEnvironment: {
93 defaultContainerCategory: "DEFAULT_CONTAINER_CATEGORY_COMPUTER_USE",
94 },
95 },
96 );
97});
98
99/** Sandbox probes, declared only when their resources are deployed. */
100const sandboxProbes = Effect.gen(function* () {
101 if (!runSandboxes) return {};
102 const getSandbox = yield* GCP.AIPlatform.GetSandboxEnvironment(Code);
103 const getTemplate =
104 yield* GCP.AIPlatform.GetSandboxEnvironmentTemplate(Browser);
105 const probes = {
106 getSandboxEnvironment: getSandbox(),
107 getSandboxEnvironmentTemplate: getTemplate(),
108 };
109 if (!runSandboxPause) return probes;
110 const pause = yield* GCP.AIPlatform.PauseSandboxEnvironment(Paused);
111 const pauseResumed = yield* GCP.AIPlatform.PauseSandboxEnvironment(Resumed);
112 const resume = yield* GCP.AIPlatform.ResumeSandboxEnvironment(Resumed);
113 return {
114 ...probes,
115 pauseSandboxEnvironment: pause({ body: {} }),
116 resumeSandboxEnvironment: Effect.gen(function* () {
117 yield* pauseResumed({ body: {} });
118 return yield* resume({ body: {} });
119 }),
120 };
121});
122
123/**
124 * Effect-native Cloud Run service exercising every AI Platform binding as
125 * its own runtime service account. Deployed from
126 * {@link ../Bindings.test.ts}.
127 */
128export default class AIPlatformBindingsHost extends GCP.Function<AIPlatformBindingsHost>()(
129 "AIPlatformBindingsHost",
130 {
131 main: import.meta.url,
132 invokerIamDisabled: true,
133 env: {
134 AIPLATFORM_SANDBOXES: runSandboxes ? "1" : "0",
135 AIPLATFORM_SANDBOX_PAUSE: runSandboxPause ? "1" : "0",
136 },
137 },
138 Effect.gen(function* () {
139 const gemini = yield* GCP.AIPlatform.GenerateContent("gemini-2.5-flash");
140 const regional = yield* GCP.AIPlatform.GenerateContent({
141 model: "gemini-2.5-flash",
142 location: "us-central1",
143 });
144 const getEngine = yield* GCP.AIPlatform.GetReasoningEngine(Agent);
145 const queryEngine = yield* GCP.AIPlatform.QueryReasoningEngine(Agent);
146 const getPipeline = yield* GCP.AIPlatform.GetTrainingPipeline(Train);
147 const cancelPipeline = yield* GCP.AIPlatform.CancelTrainingPipeline(Train);
148 const sandboxes = yield* sandboxProbes;
149
150 return {
151 fetch: serveProbes({
152 generateContent: Effect.gen(function* () {
153 const text = yield* gemini.text("Reply with exactly: pong");
154 const response = yield* regional.generate({
155 contents: [
156 { role: "user", parts: [{ text: "Reply with exactly: ping" }] },
157 ],
158 });
159 return {
160 text,
161 regional:
162 response.candidates?.[0]?.content?.parts
163 ?.map((part) => part.text ?? "")
164 .join("") ?? "",
165 modelVersion: response.modelVersion,
166 };
167 }),
168 getReasoningEngine: getEngine(),
169 queryReasoningEngine: queryEngine({
170 body: { input: { input: "hello" } },
171 }),
172 getTrainingPipeline: getPipeline(),
173 cancelTrainingPipeline: cancelPipeline({ body: {} }),
174 ...sandboxes,
175 }),
176 };
177 }).pipe(
178 Effect.provide(GCP.AIPlatform.GenerateContentHttp),
179 Effect.provide(GCP.AIPlatform.GetReasoningEngineHttp),
180 Effect.provide(GCP.AIPlatform.QueryReasoningEngineHttp),
181 Effect.provide(GCP.AIPlatform.GetTrainingPipelineHttp),
182 Effect.provide(GCP.AIPlatform.CancelTrainingPipelineHttp),
183 Effect.provide(GCP.AIPlatform.GetSandboxEnvironmentHttp),
184 Effect.provide(GCP.AIPlatform.GetSandboxEnvironmentTemplateHttp),
185 Effect.provide(GCP.AIPlatform.PauseSandboxEnvironmentHttp),
186 Effect.provide(GCP.AIPlatform.ResumeSandboxEnvironmentHttp),
187 ),
188) {}