1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Schedule from "effect/Schedule";
5import * as Stream from "effect/Stream";
6import { Unowned } from "../../AdoptPolicy.ts";
7import { isResolved } from "../../Diff.ts";
8import * as Provider from "../../Provider.ts";
9import { Resource } from "../../Resource.ts";
10import { tagRecord } from "../../Tags.ts";
11import { GcpEnvironment } from "../Environment.ts";
12import {
13 createInternalLabels,
14 diffLabels,
15 hasAlchemyLabels,
16 stripInternalLabels,
17 toLabels,
18} from "../Labels.ts";
19import type { Providers } from "../Providers.ts";
20import { listLocations } from "./names.ts";
21import { resourceNameFromOperation, waitForOperation } from "./operations.ts";
22import {
23 lastSegment,
24 locationOf,
25 locationParent,
26 parentOf,
27 toResourceId,
28} from "./ownership.ts";
29
30export type ReasoningEnginesSessionProps = {
31 /**
32 * Parent ReasoningEngine resource name
33 * (`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`).
34 * Immutable — changing it replaces the session.
35 */
36 parent: string;
37 /**
38 * Session id (the `{session}` segment). If omitted, a unique id is
39 * generated. Immutable.
40 */
41 sessionId?: string;
42 /**
43 * Immutable user id that owns the session.
44 */
45 userId: string;
46 /**
47 * User-facing display name.
48 */
49 displayName?: string;
50 /**
51 * User labels. Alchemy ownership labels are merged in automatically.
52 */
53 labels?: Record<string, string>;
54 /**
55 * Session TTL (for example `"86400s"`). Minimum 24 hours.
56 */
57 ttl?: string;
58 /**
59 * RFC3339 expiration timestamp. Minimum 24 hours from create time.
60 */
61 expireTime?: string;
62 /**
63 * Session-specific memory storing key conversation points.
64 */
65 sessionState?: Record<string, unknown>;
66};
67
68export type ReasoningEnginesSession = Resource<
69 "GCP.AIPlatform.ReasoningEnginesSession",
70 ReasoningEnginesSessionProps,
71 {
72 /** Full resource name. */
73 name: string;
74 /** Session id (last path segment). */
75 sessionId: string;
76 /** Parent ReasoningEngine resource name. */
77 parent: string;
78 /** Region id. */
79 location: string;
80 /** Project id. */
81 project: string;
82 /** User id that owns the session. */
83 userId: string | undefined;
84 /** User-facing display name. */
85 displayName: string | undefined;
86 /** User labels (Alchemy ownership labels stripped). */
87 labels: Record<string, string>;
88 /** Expiration timestamp. */
89 expireTime: string | undefined;
90 /** Session memory. */
91 sessionState: Record<string, unknown> | undefined;
92 /** RFC3339 creation timestamp. */
93 createTime: string | undefined;
94 /** RFC3339 last-update timestamp. */
95 updateTime: string | undefined;
96 },
97 never,
98 Providers
99>;
100
101/**
102 * A Vertex AI Reasoning Engine session between a user and an agent.
103 *
104 * Parent engine, session id, and user id are immutable. Display name,
105 * labels, expiration, and session state update in place.
106 *
107 * ### Creating a Session
108 * **Example:** Session for a user
109 * ```typescript
110 * const session = yield* GCP.AIPlatform.ReasoningEnginesSession("Chat", {
111 * parent: engine.name,
112 * userId: "user-123",
113 * displayName: "support-chat",
114 * });
115 * ```
116 *
117 * @resource
118 * @category AIPlatform
119 */
120export const ReasoningEnginesSession = Resource<ReasoningEnginesSession>(
121 "GCP.AIPlatform.ReasoningEnginesSession",
122);
123
124export class ReasoningEnginesSessionNotResolved extends Data.TaggedError(
125 "GCP.AIPlatform.ReasoningEnginesSessionNotResolved",
126)<{
127 name: string;
128}> {}
129
130const userLabels = (
131 labels: Record<string, string | undefined> | null | undefined,
132): Record<string, string> => stripInternalLabels(tagRecord(labels));
133
134const resourceName = (parent: string, sessionId: string) =>
135 `${parent}/sessions/${sessionId}`;
136
137const toAttrs = (
138 session: aiplatform.GoogleCloudAiplatformV1Session,
139 project: string,
140) => {
141 const name = session.name ?? "";
142 return {
143 name,
144 sessionId: lastSegment(name),
145 parent: parentOf(name),
146 location: locationOf(name),
147 project,
148 userId: session.userId,
149 displayName: session.displayName,
150 labels: userLabels(session.labels),
151 expireTime: session.expireTime,
152 sessionState: session.sessionState,
153 createTime: session.createTime,
154 updateTime: session.updateTime,
155 };
156};
157
158const getByName = (name: string) =>
159 name.length === 0
160 ? Effect.succeed(undefined)
161 : aiplatform
162 .getProjectsLocationsReasoningEnginesSessions({ name })
163 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
164
165const waitUntilExists = (name: string) =>
166 getByName(name).pipe(
167 Effect.flatMap((session) =>
168 session
169 ? Effect.succeed(session)
170 : Effect.fail(new ReasoningEnginesSessionNotResolved({ name })),
171 ),
172 Effect.retry({
173 while: (error) =>
174 error._tag === "GCP.AIPlatform.ReasoningEnginesSessionNotResolved",
175 times: 8,
176 schedule: Schedule.spaced("2 seconds"),
177 }),
178 );
179
180const listEngines = (project: string, location: string) =>
181 aiplatform.listProjectsLocationsReasoningEngines
182 .pages({
183 parent: locationParent(project, location),
184 pageSize: 100,
185 })
186 .pipe(
187 Stream.flatMap((page) =>
188 Stream.fromIterable(page.reasoningEngines ?? []),
189 ),
190 Stream.map((engine) => engine.name ?? ""),
191 Stream.filter((name) => name.length > 0),
192 Stream.runCollect,
193 Effect.map((chunk) => Array.from(chunk)),
194 Effect.catchTag("NotFound", () => Effect.succeed([] as string[])),
195 );
196
197const listAtParent = (parent: string, project: string) =>
198 aiplatform.listProjectsLocationsReasoningEnginesSessions
199 .pages({ parent, pageSize: 100 })
200 .pipe(
201 Stream.flatMap((page) => Stream.fromIterable(page.sessions ?? [])),
202 Stream.filter((session) =>
203 Object.keys(session.labels ?? {}).some((key) =>
204 key.startsWith("alchemy-"),
205 ),
206 ),
207 Stream.map((session) => toAttrs(session, project)),
208 Stream.runCollect,
209 Effect.map((chunk) => Array.from(chunk)),
210 Effect.catchTag("NotFound", () => Effect.succeed([])),
211 );
212
213export const ReasoningEnginesSessionProvider = () =>
214 Provider.succeed(ReasoningEnginesSession, {
215 stables: [
216 "name",
217 "sessionId",
218 "parent",
219 "location",
220 "project",
221 "userId",
222 "createTime",
223 ],
224
225 diff: Effect.fn(function* ({ news, olds, output }) {
226 if (!isResolved(news)) return undefined;
227 const previousParent = olds?.parent ?? output?.parent;
228 if (previousParent !== undefined && news.parent !== previousParent) {
229 return { action: "replace" as const, deleteFirst: false };
230 }
231 const previousId = olds?.sessionId ?? output?.sessionId;
232 if (
233 previousId !== undefined &&
234 news.sessionId !== undefined &&
235 news.sessionId !== previousId
236 ) {
237 return { action: "replace" as const, deleteFirst: true };
238 }
239 const previousUser = olds?.userId ?? output?.userId;
240 if (previousUser !== undefined && news.userId !== previousUser) {
241 return { action: "replace" as const, deleteFirst: true };
242 }
243 return undefined;
244 }),
245
246 read: Effect.fn(function* ({ id, olds, output }) {
247 const env = yield* GcpEnvironment.current;
248 const sessionId = yield* toResourceId(
249 id,
250 olds?.sessionId,
251 output?.sessionId,
252 );
253 const name =
254 output?.name ??
255 (olds?.parent !== undefined
256 ? resourceName(olds.parent, sessionId)
257 : "");
258 const existing = yield* getByName(name);
259 if (existing === undefined) return undefined;
260 const attrs = toAttrs(existing, env.project);
261 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
262 ? attrs
263 : Unowned(attrs);
264 }),
265
266 list: () =>
267 Effect.gen(function* () {
268 const env = yield* GcpEnvironment.current;
269 const engines = (yield* Effect.forEach(
270 listLocations(env.region),
271 (location) => listEngines(env.project, location),
272 )).flat();
273 const pages = yield* Effect.forEach(
274 engines,
275 (parent) => listAtParent(parent, env.project),
276 { concurrency: 4 },
277 );
278 return pages.flat();
279 }),
280
281 reconcile: Effect.fn(function* ({ id, news, output }) {
282 const env = yield* GcpEnvironment.current;
283 const sessionId = yield* toResourceId(
284 id,
285 news.sessionId,
286 output?.sessionId,
287 );
288 const name = resourceName(news.parent, sessionId);
289 const desiredLabels = {
290 ...toLabels(news.labels),
291 ...(yield* createInternalLabels(id)),
292 };
293
294 let current = yield* getByName(output?.name ?? name);
295
296 if (current === undefined) {
297 const created = yield* aiplatform
298 .createProjectsLocationsReasoningEnginesSessions({
299 parent: news.parent,
300 sessionId,
301 body: {
302 userId: news.userId,
303 displayName: news.displayName,
304 labels: desiredLabels,
305 ttl: news.ttl,
306 expireTime: news.expireTime,
307 sessionState: news.sessionState,
308 },
309 })
310 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
311 if (created !== undefined) {
312 const done = yield* waitForOperation(created, {
313 alreadyExistsOk: true,
314 });
315 const createdName = resourceNameFromOperation(done) ?? name;
316 current = yield* waitUntilExists(createdName);
317 } else {
318 current = yield* getByName(name);
319 }
320 }
321
322 if (current === undefined) {
323 return yield* new ReasoningEnginesSessionNotResolved({ name });
324 }
325
326 const observedLabels = tagRecord(current.labels);
327 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
328 const labelsChanged = upsert.length > 0 || removed.length > 0;
329 const displayChanged =
330 (current.displayName ?? "") !== (news.displayName ?? "");
331 const expireChanged =
332 news.expireTime !== undefined &&
333 (current.expireTime ?? "") !== news.expireTime;
334 const stateChanged =
335 news.sessionState !== undefined &&
336 JSON.stringify(current.sessionState ?? {}) !==
337 JSON.stringify(news.sessionState ?? {});
338
339 if (labelsChanged || displayChanged || expireChanged || stateChanged) {
340 current =
341 yield* aiplatform.patchProjectsLocationsReasoningEnginesSessions({
342 name: current.name ?? name,
343 updateMask: [
344 labelsChanged ? "labels" : undefined,
345 displayChanged ? "displayName" : undefined,
346 expireChanged ? "expireTime" : undefined,
347 stateChanged ? "sessionState" : undefined,
348 ]
349 .filter((field): field is string => field !== undefined)
350 .join(","),
351 body: {
352 name: current.name ?? name,
353 displayName: news.displayName,
354 labels: desiredLabels,
355 expireTime: news.expireTime,
356 sessionState: news.sessionState,
357 },
358 });
359 }
360
361 return toAttrs(current, env.project);
362 }),
363
364 delete: Effect.fn(function* ({ output }) {
365 const existing = yield* getByName(output.name);
366 if (existing === undefined) return;
367 const operation = yield* aiplatform
368 .deleteProjectsLocationsReasoningEnginesSessions({
369 name: output.name,
370 })
371 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
372 if (operation !== undefined) {
373 yield* waitForOperation(operation, { notFoundOk: true });
374 }
375 }),
376 });