feat(go): add DefineSessionFlow and DefineSessionFlowFromPrompt

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

@@ -0,0 +1,120 @@
+/**
+ * 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 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(),
+  /** ID of a snapshot that was just persisted. */
+  snapshotId: z.string().optional(),
+  /** Signals that the session flow has finished processing the current input. */
+  endTurn: z.boolean().optional(),
+});
+export type SessionFlowStreamChunk = z.infer<
+  typeof SessionFlowStreamChunkSchema
+>;

genkit-tools/genkit-schema.json

@@ -1,6 +1,140 @@
 {
   "$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"
+        },
+        "snapshotId": {
+          "type": "string"
+        },
+        "endTurn": {
+          "type": "boolean"
+        }
+      },
+      "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"
+      ]
+    },
     "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',

go/ai/exp/gen.go

@@ -0,0 +1,122 @@
+// 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.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+// This file was generated by jsonschemagen. DO NOT EDIT.
+
+package exp
+
+import (
+	"github.com/firebase/genkit/go/ai"
+)
+
+// SessionFlowInit is the input for starting an session flow invocation.
+// Provide either SnapshotID (to load from store) or State (direct state).
+type SessionFlowInit[State any] struct {
+	// SnapshotID loads state from a persisted snapshot.
+	// Mutually exclusive with State.
+	SnapshotID string `json:"snapshotId,omitempty"`
+	// State provides direct state for the invocation.
+	// Mutually exclusive with SnapshotID.
+	State *SessionState[State] `json:"state,omitempty"`
+}
+
+// SessionFlowInput is the input sent to an session flow during a conversation turn.
+type SessionFlowInput struct {
+	// Messages contains the user's input for this turn.
+	Messages []*ai.Message `json:"messages,omitempty"`
+	// ToolRestarts contains tool request parts to re-execute interrupted tools.
+	// Use [ai.ToolDef.RestartWith] to create these parts from an interrupted
+	// tool request. When set, the generate call resumes with these restarts
+	// instead of treating Messages as tool responses.
+	ToolRestarts []*ai.Part `json:"toolRestarts,omitempty"`
+}
+
+// SessionFlowOutput is the output when an session flow invocation completes.
+// It wraps SessionFlowResult with framework-managed fields.
+type SessionFlowOutput[State any] struct {
+	// Artifacts contains artifacts produced during the session.
+	Artifacts []*Artifact `json:"artifacts,omitempty"`
+	// Message is the last model response message from the conversation.
+	Message *ai.Message `json:"message,omitempty"`
+	// SnapshotID is the ID of the snapshot created at the end of this invocation.
+	// Empty if no snapshot was created (callback returned false or no store configured).
+	SnapshotID string `json:"snapshotId,omitempty"`
+	// State contains the final conversation state.
+	// Only populated when state is client-managed (no store configured).
+	State *SessionState[State] `json:"state,omitempty"`
+}
+
+// SessionFlowResult is the return value from an SessionFlowFunc.
+// It contains the user-specified outputs of the agent invocation.
+type SessionFlowResult struct {
+	// Artifacts contains artifacts produced during the session.
+	Artifacts []*Artifact `json:"artifacts,omitempty"`
+	// Message is the last model response message from the conversation.
+	Message *ai.Message `json:"message,omitempty"`
+}
+
+// SessionFlowStreamChunk represents a single item in the session flow's output stream.
+// Multiple fields can be populated in a single chunk.
+type SessionFlowStreamChunk[Stream any] struct {
+	// Artifact contains a newly produced artifact.
+	Artifact *Artifact `json:"artifact,omitempty"`
+	// EndTurn signals that the session flow has finished processing the current input.
+	// When true, the client should stop iterating and may send the next input.
+	EndTurn bool `json:"endTurn,omitempty"`
+	// ModelChunk contains generation tokens from the model.
+	ModelChunk *ai.ModelResponseChunk `json:"modelChunk,omitempty"`
+	// SnapshotID contains the ID of a snapshot that was just persisted.
+	SnapshotID string `json:"snapshotId,omitempty"`
+	// Status contains user-defined structured status information.
+	// The Stream type parameter defines the shape of this data.
+	Status Stream `json:"status,omitempty"`
+}
+
+// Artifact represents a named collection of parts produced during a session.
+// Examples: generated files, images, code snippets, diagrams, etc.
+type Artifact struct {
+	// Metadata contains additional artifact-specific data.
+	Metadata map[string]any `json:"metadata,omitempty"`
+	// Name identifies the artifact (e.g., "generated_code.go", "diagram.png").
+	Name string `json:"name,omitempty"`
+	// Parts contains the artifact content (text, media, etc.).
+	Parts []*ai.Part `json:"parts"`
+}
+
+// SessionState is the portable conversation state that flows between client
+// and server. It contains only the data needed for conversation continuity.
+type SessionState[State any] struct {
+	// Artifacts are named collections of parts produced during the conversation.
+	Artifacts []*Artifact `json:"artifacts,omitempty"`
+	// Custom is the user-defined state associated with this conversation.
+	Custom State `json:"custom,omitempty"`
+	// InputVariables is the input used for session flows that require input variables
+	// (e.g. prompt-backed session flows).
+	InputVariables any `json:"inputVariables,omitempty"`
+	// Messages is the conversation history (user/model exchanges).
+	// Does NOT include prompt-rendered messages — those are rendered fresh each turn.
+	Messages []*ai.Message `json:"messages,omitempty"`
+}
+
+// SnapshotEvent identifies what triggered a snapshot.
+type SnapshotEvent string
+
+const (
+	// TurnEnd indicates the snapshot was triggered at the end of a turn.
+	SnapshotEventTurnEnd SnapshotEvent = "turnEnd"
+	// InvocationEnd indicates the snapshot was triggered at the end of the invocation.
+	SnapshotEventInvocationEnd SnapshotEvent = "invocationEnd"
+)