feat(go): add DefineSessionFlow and DefineSessionFlowFromPrompt

genkit-tools/common/src/types/agent.ts

@@ -0,0 +1,139 @@
+/**
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { z } from 'zod';
+import { MessageSchema, ModelResponseChunkSchema } from './model';
+import { PartSchema } from './parts';
+
+/**
+ * Zod schema for an artifact produced during a session.
+ */
+export const ArtifactSchema = z.object({
+  /** Name identifies the artifact (e.g., "generated_code.go", "diagram.png"). */
+  name: z.string().optional(),
+  /** Parts contains the artifact content (text, media, etc.). */
+  parts: z.array(PartSchema),
+  /** Metadata contains additional artifact-specific data. */
+  metadata: z.record(z.any()).optional(),
+});
+export type Artifact = z.infer<typeof ArtifactSchema>;
+
+/**
+ * Zod schema for snapshot event.
+ */
+export const SnapshotEventSchema = z.enum(['turnEnd', 'invocationEnd']);
+export type SnapshotEvent = z.infer<typeof SnapshotEventSchema>;
+
+/**
+ * Zod schema for session state.
+ */
+export const SessionStateSchema = z.object({
+  /** Conversation history (user/model exchanges). */
+  messages: z.array(MessageSchema).optional(),
+  /** User-defined state associated with this conversation. */
+  custom: z.any().optional(),
+  /** Named collections of parts produced during the conversation. */
+  artifacts: z.array(ArtifactSchema).optional(),
+  /** Input used for session flows that require input variables. */
+  inputVariables: z.any().optional(),
+});
+export type SessionState = z.infer<typeof SessionStateSchema>;
+
+/**
+ * Zod schema for session flow input (per-turn).
+ */
+export const SessionFlowInputSchema = z.object({
+  /** User's input messages for this turn. */
+  messages: z.array(MessageSchema).optional(),
+  /** Tool request parts to re-execute interrupted tools. */
+  toolRestarts: z.array(PartSchema).optional(),
+});
+export type SessionFlowInput = z.infer<typeof SessionFlowInputSchema>;
+
+/**
+ * Zod schema for session flow initialization.
+ */
+export const SessionFlowInitSchema = z.object({
+  /** Loads state from a persisted snapshot. Mutually exclusive with state. */
+  snapshotId: z.string().optional(),
+  /** Direct state for the invocation. Mutually exclusive with snapshotId. */
+  state: SessionStateSchema.optional(),
+});
+export type SessionFlowInit = z.infer<typeof SessionFlowInitSchema>;
+
+/**
+ * Zod schema for session flow result.
+ */
+export const SessionFlowResultSchema = z.object({
+  /** Last model response message from the conversation. */
+  message: MessageSchema.optional(),
+  /** Artifacts produced during the session. */
+  artifacts: z.array(ArtifactSchema).optional(),
+});
+export type SessionFlowResult = z.infer<typeof SessionFlowResultSchema>;
+
+/**
+ * Zod schema for session flow output.
+ */
+export const SessionFlowOutputSchema = z.object({
+  /** ID of the snapshot created at the end of this invocation. */
+  snapshotId: z.string().optional(),
+  /** Final conversation state (only when client-managed). */
+  state: SessionStateSchema.optional(),
+  /** Last model response message from the conversation. */
+  message: MessageSchema.optional(),
+  /** Artifacts produced during the session. */
+  artifacts: z.array(ArtifactSchema).optional(),
+});
+export type SessionFlowOutput = z.infer<typeof SessionFlowOutputSchema>;
+
+/**
+ * Zod schema for the turn-end signal emitted by a session flow.
+ *
+ * A TurnEnd value is emitted exactly once per turn, regardless of whether a
+ * snapshot was persisted. Grouping all turn-end signals here lets callers
+ * detect turn boundaries with a single field check and leaves room for
+ * additional turn-end metadata in the future.
+ */
+export const TurnEndSchema = z.object({
+  /**
+   * ID of the snapshot persisted at the end of this turn. Empty if no
+   * snapshot was created (callback returned false or no store configured).
+   */
+  snapshotId: z.string().optional(),
+});
+export type TurnEnd = z.infer<typeof TurnEndSchema>;
+
+/**
+ * Zod schema for session flow stream chunk.
+ */
+export const SessionFlowStreamChunkSchema = z.object({
+  /** Generation tokens from the model. */
+  modelChunk: ModelResponseChunkSchema.optional(),
+  /** User-defined structured status information. */
+  status: z.any().optional(),
+  /** A newly produced artifact. */
+  artifact: ArtifactSchema.optional(),
+  /**
+   * Non-null when the session flow has finished processing the current
+   * input. Groups all turn-end signals; the client should stop iterating and
+   * may send the next input.
+   */
+  turnEnd: TurnEndSchema.optional(),
+});
+export type SessionFlowStreamChunk = z.infer<
+  typeof SessionFlowStreamChunkSchema
+>;

genkit-tools/genkit-schema.json

@@ -1,6 +1,146 @@
 {
   "$schema": "http://json-schema.org/draft-07/schema#",
   "$defs": {
+    "Artifact": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string"
+        },
+        "parts": {
+          "type": "array",
+          "items": {
+            "$ref": "#/$defs/Part"
+          }
+        },
+        "metadata": {
+          "type": "object",
+          "additionalProperties": {}
+        }
+      },
+      "required": [
+        "parts"
+      ],
+      "additionalProperties": false
+    },
+    "SessionFlowInit": {
+      "type": "object",
+      "properties": {
+        "snapshotId": {
+          "type": "string"
+        },
+        "state": {
+          "$ref": "#/$defs/SessionState"
+        }
+      },
+      "additionalProperties": false
+    },
+    "SessionFlowInput": {
+      "type": "object",
+      "properties": {
+        "messages": {
+          "type": "array",
+          "items": {
+            "$ref": "#/$defs/Message"
+          }
+        },
+        "toolRestarts": {
+          "type": "array",
+          "items": {
+            "$ref": "#/$defs/Part"
+          }
+        }
+      },
+      "additionalProperties": false
+    },
+    "SessionFlowOutput": {
+      "type": "object",
+      "properties": {
+        "snapshotId": {
+          "type": "string"
+        },
+        "state": {
+          "$ref": "#/$defs/SessionState"
+        },
+        "message": {
+          "$ref": "#/$defs/Message"
+        },
+        "artifacts": {
+          "type": "array",
+          "items": {
+            "$ref": "#/$defs/Artifact"
+          }
+        }
+      },
+      "additionalProperties": false
+    },
+    "SessionFlowResult": {
+      "type": "object",
+      "properties": {
+        "message": {
+          "$ref": "#/$defs/Message"
+        },
+        "artifacts": {
+          "type": "array",
+          "items": {
+            "$ref": "#/$defs/Artifact"
+          }
+        }
+      },
+      "additionalProperties": false
+    },
+    "SessionFlowStreamChunk": {
+      "type": "object",
+      "properties": {
+        "modelChunk": {
+          "$ref": "#/$defs/ModelResponseChunk"
+        },
+        "status": {},
+        "artifact": {
+          "$ref": "#/$defs/Artifact"
+        },
+        "turnEnd": {
+          "$ref": "#/$defs/TurnEnd"
+        }
+      },
+      "additionalProperties": false
+    },
+    "SessionState": {
+      "type": "object",
+      "properties": {
+        "messages": {
+          "type": "array",
+          "items": {
+            "$ref": "#/$defs/Message"
+          }
+        },
+        "custom": {},
+        "artifacts": {
+          "type": "array",
+          "items": {
+            "$ref": "#/$defs/Artifact"
+          }
+        },
+        "inputVariables": {}
+      },
+      "additionalProperties": false
+    },
+    "SnapshotEvent": {
+      "type": "string",
+      "enum": [
+        "turnEnd",
+        "invocationEnd"
+      ]
+    },
+    "TurnEnd": {
+      "type": "object",
+      "properties": {
+        "snapshotId": {
+          "type": "string"
+        }
+      },
+      "additionalProperties": false
+    },
     "DocumentData": {
       "type": "object",
       "properties": {

genkit-tools/scripts/schema-exporter.ts

@@ -22,6 +22,7 @@ import { zodToJsonSchema } from 'zod-to-json-schema';
 
 /** List of files that contain types to be exported. */
 const EXPORTED_TYPE_MODULES = [
+  '../common/src/types/agent.ts',
   '../common/src/types/document.ts',
   '../common/src/types/embedder.ts',
   '../common/src/types/evaluator.ts',