From 42d13ec1cb916da156f4e400fe0751262134c54e Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 8 Dec 2025 18:47:45 +0800 Subject: [PATCH] Refactor chat storage design and enhance data models - Updated the chat storage design to reflect a shift from "Conversation" to "Chat" terminology, improving clarity in the data structure. - Revised the data models for Chat and Message tables, including new fields for enhanced metadata management and permissions. - Introduced a `space_snapshot` field in the Step model to facilitate recovery during execution, improving resilience in chat interactions. - Enhanced middleware documentation to reflect the modular architecture and provide clearer usage examples for different API routes. - Updated tests to ensure the integrity of new data structures and functionalities, reinforcing the robustness of the chat storage system. --- agent/store/CHAT_STORAGE_DESIGN.md | 401 +++++++++++------- openapi/request/REQUEST_DESIGN.md | 636 +++++++++++++++++++++-------- 2 files changed, 716 insertions(+), 321 deletions(-) diff --git a/agent/store/CHAT_STORAGE_DESIGN.md b/agent/store/CHAT_STORAGE_DESIGN.md index 5c67750f..69f22ca8 100644 --- a/agent/store/CHAT_STORAGE_DESIGN.md +++ b/agent/store/CHAT_STORAGE_DESIGN.md @@ -39,14 +39,14 @@ The chat storage system is designed to: The Agent storage focuses on **chat content and execution state**, while request tracking (billing, rate limiting, auditing) is handled globally by the OpenAPI layer: -| Concern | Module | Table | -| ------------------ | ----------------- | -------------------- | -| Request tracking | `openapi/request` | `openapi_request` | -| Billing (tokens) | `openapi/request` | `openapi_request` | -| Rate limiting | `openapi/request` | - | -| Chat conversations | `agent/store` | `agent_conversation` | -| Chat messages | `agent/store` | `agent_message` | -| Execution steps | `agent/store` | `agent_step` | +| Concern | Module | Table | +| ---------------- | ----------------- | ----------------- | +| Request tracking | `openapi/request` | `openapi_request` | +| Billing (tokens) | `openapi/request` | `openapi_request` | +| Rate limiting | `openapi/request` | - | +| Chat sessions | `agent/store` | `agent_chat` | +| Chat messages | `agent/store` | `agent_message` | +| Execution steps | `agent/store` | `agent_step` | The `request_id` from OpenAPI middleware is passed to Agent and stored in messages/steps for correlation. @@ -58,7 +58,7 @@ The `request_id` from OpenAPI middleware is passed to Agent and stored in messag ├─────────────────────────────────────────────────────────────┤ │ │ │ ┌─────────────────┐ │ -│ │ Conversation │ Metadata: title, assistant, user │ +│ │ Chat │ Metadata: title, assistant, user │ │ └────────┬────────┘ │ │ │ │ │ │ 1:N │ @@ -78,35 +78,60 @@ The `request_id` from OpenAPI middleware is passed to Agent and stored in messag ## Data Models -### 1. Conversation Table +### 1. Chat Table -Stores conversation metadata and session information. +Stores chat metadata and session information. -**Table Name:** `agent_conversation` +**Table Name:** `agent_chat` -| Column | Type | Nullable | Index | Description | -| ----------------- | ----------- | -------- | ------ | ----------------------------------- | -| `id` | ID | No | PK | Auto-increment primary key | -| `conversation_id` | string(64) | No | Unique | Unique conversation identifier | -| `title` | string(500) | Yes | - | Conversation title | -| `assistant_id` | string(200) | No | Yes | Associated assistant ID | -| `user_id` | string(200) | No | Yes | Owner user ID | -| `team_id` | string(200) | Yes | Yes | Team ID for access control | -| `mode` | string(50) | No | - | Conversation mode (default: "chat") | -| `status` | enum | No | Yes | Status: `active`, `archived` | -| `last_message_at` | timestamp | Yes | Yes | Timestamp of last message | -| `metadata` | json | Yes | - | Additional metadata | -| `created_at` | timestamp | No | Yes | Creation timestamp | -| `updated_at` | timestamp | No | - | Last update timestamp | +| Column | Type | Nullable | Index | Description | +| ----------------- | ----------- | -------- | ------ | -------------------------------- | +| `id` | ID | No | PK | Auto-increment primary key | +| `chat_id` | string(64) | No | Unique | Unique chat identifier | +| `title` | string(500) | Yes | - | Chat title | +| `assistant_id` | string(200) | No | Yes | Associated assistant ID | +| `mode` | string(50) | No | - | Chat mode (default: "chat") | +| `status` | enum | No | Yes | Status: `active`, `archived` | +| `preset` | boolean | No | - | Whether this is a preset chat | +| `public` | boolean | No | - | Whether shared across all teams | +| `share` | enum | No | Yes | Sharing scope: `private`, `team` | +| `sort` | integer | No | - | Sort order for display | +| `last_message_at` | timestamp | Yes | Yes | Timestamp of last message | +| `metadata` | json | Yes | - | Additional metadata | +| `created_at` | timestamp | No | Yes | Creation timestamp | +| `updated_at` | timestamp | No | - | Last update timestamp | + +**Model Options:** + +```json +{ + "option": { + "soft_deletes": true, + "permission": true, + "timestamps": true + } +} +``` + +**Note:** `permission: true` enables Yao's built-in permission management, which automatically adds the following fields: + +| Field | Type | Description | +| ------------------ | ----------- | ------------------------------ | +| `__yao_created_by` | string(200) | User ID who created the record | +| `__yao_updated_by` | string(200) | User ID who last updated | +| `__yao_team_id` | string(200) | Team ID for team-level access | +| `__yao_tenant_id` | string(200) | Tenant ID for multi-tenancy | + +These fields are automatically managed by the framework and used for access control filtering. **Indexes:** -| Name | Columns | Type | -| -------------------- | ------------------- | ----- | -| `idx_conv_user` | `user_id`, `status` | index | -| `idx_conv_team` | `team_id`, `status` | index | -| `idx_conv_assistant` | `assistant_id` | index | -| `idx_conv_last_msg` | `last_message_at` | index | +| Name | Columns | Type | +| -------------------- | ----------------- | ----- | +| `idx_chat_assistant` | `assistant_id` | index | +| `idx_chat_status` | `status` | index | +| `idx_chat_share` | `share` | index | +| `idx_chat_last_msg` | `last_message_at` | index | ### 2. Message Table @@ -114,31 +139,31 @@ Stores user-visible messages (both user input and assistant responses). **Table Name:** `agent_message` -| Column | Type | Nullable | Index | Description | -| ----------------- | ----------- | -------- | ------ | ----------------------------------------- | -| `id` | ID | No | PK | Auto-increment primary key | -| `message_id` | string(64) | No | Unique | Unique message identifier | -| `conversation_id` | string(64) | No | Yes | Parent conversation ID | -| `request_id` | string(64) | Yes | Yes | Request ID for grouping | -| `role` | enum | No | Yes | Role: `user`, `assistant` | -| `type` | string(50) | No | - | Message type (text, image, loading, etc.) | -| `props` | json | No | - | Message properties (content, url, etc.) | -| `block_id` | string(64) | Yes | Yes | Block grouping ID | -| `thread_id` | string(64) | Yes | Yes | Thread grouping ID | -| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) | -| `sequence` | integer | No | Yes | Message order within conversation | -| `metadata` | json | Yes | - | Additional metadata | -| `created_at` | timestamp | No | Yes | Creation timestamp | -| `updated_at` | timestamp | No | - | Last update timestamp | +| Column | Type | Nullable | Index | Description | +| -------------- | ----------- | -------- | ------ | ----------------------------------------- | +| `id` | ID | No | PK | Auto-increment primary key | +| `message_id` | string(64) | No | Unique | Unique message identifier | +| `chat_id` | string(64) | No | Yes | Parent chat ID | +| `request_id` | string(64) | Yes | Yes | Request ID for grouping | +| `role` | enum | No | Yes | Role: `user`, `assistant` | +| `type` | string(50) | No | - | Message type (text, image, loading, etc.) | +| `props` | json | No | - | Message properties (content, url, etc.) | +| `block_id` | string(64) | Yes | Yes | Block grouping ID | +| `thread_id` | string(64) | Yes | Yes | Thread grouping ID | +| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) | +| `sequence` | integer | No | Yes | Message order within chat | +| `metadata` | json | Yes | - | Additional metadata | +| `created_at` | timestamp | No | Yes | Creation timestamp | +| `updated_at` | timestamp | No | - | Last update timestamp | **Indexes:** -| Name | Columns | Type | -| ------------------- | ----------------------------- | ----- | -| `idx_msg_conv_seq` | `conversation_id`, `sequence` | index | -| `idx_msg_request` | `request_id` | index | -| `idx_msg_block` | `block_id` | index | -| `idx_msg_assistant` | `assistant_id` | index | +| Name | Columns | Type | +| ------------------- | --------------------- | ----- | +| `idx_msg_chat_seq` | `chat_id`, `sequence` | index | +| `idx_msg_request` | `request_id` | index | +| `idx_msg_block` | `block_id` | index | +| `idx_msg_assistant` | `assistant_id` | index | **Message Types:** @@ -161,7 +186,7 @@ Stores execution steps for resume/retry functionality. | ----------------- | ----------- | -------- | ------ | -------------------------------- | | `id` | ID | No | PK | Auto-increment primary key | | `step_id` | string(64) | No | Unique | Unique step identifier | -| `conversation_id` | string(64) | No | Yes | Parent conversation ID | +| `chat_id` | string(64) | No | Yes | Parent chat ID | | `request_id` | string(64) | No | Yes | Request ID | | `assistant_id` | string(200) | No | Yes | Assistant executing this step | | `stack_id` | string(64) | No | Yes | Stack node ID for this execution | @@ -171,12 +196,34 @@ Stores execution steps for resume/retry functionality. | `status` | enum | No | Yes | Step status | | `input` | json | Yes | - | Step input data | | `output` | json | Yes | - | Step output data | +| `space_snapshot` | json | Yes | - | Space data snapshot for recovery | | `error` | text | Yes | - | Error message if failed | | `sequence` | integer | No | Yes | Step order within request | | `metadata` | json | Yes | - | Additional metadata | | `created_at` | timestamp | No | Yes | Creation timestamp | | `updated_at` | timestamp | No | - | Last update timestamp | +**Space Snapshot:** + +The `space_snapshot` field stores the shared data space (`ctx.Space`) at each step for recovery purposes. + +```typescript +// Example: In Next hook, set data to Space before delegate +ctx.space.Set("choose_prompt", "query"); +return { + delegate: { agent_id: "expense", messages: payload.messages }, +}; +``` + +If interrupted during delegate, the `space_snapshot` allows restoring `ctx.Space` state: + +```json +{ + "choose_prompt": "query", + "user_preferences": { "currency": "USD" } +} +``` + **Step Types:** | Type | Description | Input | Output | @@ -202,7 +249,7 @@ Stores execution steps for resume/retry functionality. | Name | Columns | Type | | -------------------- | ------------------------ | ----- | -| `idx_step_conv` | `conversation_id` | index | +| `idx_step_chat` | `chat_id` | index | | `idx_step_request` | `request_id`, `sequence` | index | | `idx_step_status` | `status` | index | | `idx_step_stack` | `stack_id` | index | @@ -350,19 +397,26 @@ func (ast *Assistant) Stream(ctx, inputMessages, options) { // createStep creates a step with context information func createStep(ctx *Context, stepType, status string, input, output interface{}) *Step { + // Capture Space snapshot for recovery + var spaceSnapshot map[string]interface{} + if ctx.Space != nil { + spaceSnapshot = ctx.Space.Snapshot() // Get all key-value pairs + } + return &Step{ - StepID: generateID(), - ConversationID: ctx.ChatID, // ChatID = conversation_id - RequestID: ctx.RequestID, // From OpenAPI middleware - AssistantID: ctx.AssistantID, - StackID: ctx.Stack.ID, - StackParentID: ctx.Stack.ParentID, - StackDepth: ctx.Stack.Depth, - Type: stepType, - Status: status, - Input: input, - Output: output, - Sequence: nextSequence(), + StepID: generateID(), + ChatID: ctx.ChatID, // ChatID + RequestID: ctx.RequestID, // From OpenAPI middleware + AssistantID: ctx.AssistantID, + StackID: ctx.Stack.ID, + StackParentID: ctx.Stack.ParentID, + StackDepth: ctx.Stack.Depth, + Type: stepType, + Status: status, + Input: input, + Output: output, + SpaceSnapshot: spaceSnapshot, // Shared space data for recovery + Sequence: nextSequence(), } } @@ -373,88 +427,103 @@ func createStep(ctx *Context, stepType, status string, input, output interface{} ```go // ChatStore defines the chat storage interface type ChatStore interface { - // Conversation Management - CreateConversation(conv *Conversation) error - GetConversation(conversationID string) (*Conversation, error) - UpdateConversation(conversationID string, updates map[string]interface{}) error - DeleteConversation(conversationID string) error - ListConversations(filter ConversationFilter) (*ConversationList, error) + // Chat Management + CreateChat(chat *Chat) error + GetChat(chatID string) (*Chat, error) + UpdateChat(chatID string, updates map[string]interface{}) error + DeleteChat(chatID string) error + ListChats(filter ChatFilter) (*ChatList, error) // Message Management - SaveMessages(conversationID string, messages []*Message) error - GetMessages(conversationID string, filter MessageFilter) ([]*Message, error) + SaveMessages(chatID string, messages []*Message) error + GetMessages(chatID string, filter MessageFilter) ([]*Message, error) UpdateMessage(messageID string, updates map[string]interface{}) error - DeleteMessages(conversationID string, messageIDs []string) error + DeleteMessages(chatID string, messageIDs []string) error // Step Management - SaveStep(step *Step) error + SaveSteps(steps []*Step) error UpdateStep(stepID string, updates map[string]interface{}) error GetSteps(requestID string) ([]*Step, error) - GetLastIncompleteStep(conversationID string) (*Step, error) + GetLastIncompleteStep(chatID string) (*Step, error) + GetStepsByStackID(stackID string) ([]*Step, error) + GetStackPath(stackID string) ([]string, error) // Returns [root_stack_id, ..., current_stack_id] +} + +// SpaceStore defines the interface for Space snapshot operations +// Note: Space itself uses plan.Space interface, this is for persistence +type SpaceStore interface { + // Snapshot returns all key-value pairs in the space + Snapshot() map[string]interface{} + + // Restore sets multiple key-value pairs from a snapshot + Restore(data map[string]interface{}) error } ```` ### Data Structures ```go -// Conversation represents a chat conversation -type Conversation struct { - ConversationID string `json:"conversation_id"` - Title string `json:"title,omitempty"` - AssistantID string `json:"assistant_id"` - UserID string `json:"user_id"` - TeamID string `json:"team_id,omitempty"` - Mode string `json:"mode"` - Status string `json:"status"` - LastMessageAt *time.Time `json:"last_message_at,omitempty"` - Metadata map[string]interface{} `json:"metadata,omitempty"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` +// Chat represents a chat session +type Chat struct { + ChatID string `json:"chat_id"` + Title string `json:"title,omitempty"` + AssistantID string `json:"assistant_id"` + Mode string `json:"mode"` + Status string `json:"status"` + Preset bool `json:"preset"` + Public bool `json:"public"` + Share string `json:"share"` // "private" or "team" + Sort int `json:"sort"` + LastMessageAt *time.Time `json:"last_message_at,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } // Message represents a chat message type Message struct { - MessageID string `json:"message_id"` - ConversationID string `json:"conversation_id"` - RequestID string `json:"request_id,omitempty"` - Role string `json:"role"` - Type string `json:"type"` - Props map[string]interface{} `json:"props"` - BlockID string `json:"block_id,omitempty"` - ThreadID string `json:"thread_id,omitempty"` - AssistantID string `json:"assistant_id,omitempty"` - Sequence int `json:"sequence"` - Metadata map[string]interface{} `json:"metadata,omitempty"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + MessageID string `json:"message_id"` + ChatID string `json:"chat_id"` + RequestID string `json:"request_id,omitempty"` + Role string `json:"role"` + Type string `json:"type"` + Props map[string]interface{} `json:"props"` + BlockID string `json:"block_id,omitempty"` + ThreadID string `json:"thread_id,omitempty"` + AssistantID string `json:"assistant_id,omitempty"` + Sequence int `json:"sequence"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } // Step represents an execution step type Step struct { - StepID string `json:"step_id"` - ConversationID string `json:"conversation_id"` - RequestID string `json:"request_id"` - AssistantID string `json:"assistant_id"` - StackID string `json:"stack_id"` - StackParentID string `json:"stack_parent_id,omitempty"` - StackDepth int `json:"stack_depth"` - Type string `json:"type"` - Status string `json:"status"` - Input map[string]interface{} `json:"input,omitempty"` - Output map[string]interface{} `json:"output,omitempty"` - Error string `json:"error,omitempty"` - Sequence int `json:"sequence"` - Metadata map[string]interface{} `json:"metadata,omitempty"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + StepID string `json:"step_id"` + ChatID string `json:"chat_id"` + RequestID string `json:"request_id"` + AssistantID string `json:"assistant_id"` + StackID string `json:"stack_id"` + StackParentID string `json:"stack_parent_id,omitempty"` + StackDepth int `json:"stack_depth"` + Type string `json:"type"` + Status string `json:"status"` + Input map[string]interface{} `json:"input,omitempty"` + Output map[string]interface{} `json:"output,omitempty"` + SpaceSnapshot map[string]interface{} `json:"space_snapshot,omitempty"` // Shared space data for recovery + Error string `json:"error,omitempty"` + Sequence int `json:"sequence"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } ``` ### Filter Structures ```go -// ConversationFilter for listing conversations -type ConversationFilter struct { +// ChatFilter for listing chats +type ChatFilter struct { UserID string `json:"user_id,omitempty"` TeamID string `json:"team_id,omitempty"` AssistantID string `json:"assistant_id,omitempty"` @@ -473,13 +542,13 @@ type MessageFilter struct { Offset int `json:"offset,omitempty"` } -// ConversationList paginated response -type ConversationList struct { - Data []*Conversation `json:"data"` - Page int `json:"page"` - PageSize int `json:"pagesize"` - PageCount int `json:"pagecount"` - Total int `json:"total"` +// ChatList paginated response +type ChatList struct { + Data []*Chat `json:"data"` + Page int `json:"page"` + PageSize int `json:"pagesize"` + PageCount int `json:"pagecount"` + Total int `json:"total"` } ``` @@ -492,23 +561,23 @@ See [Write Strategy - Implementation](#implementation) for the complete flow wit ### 2. Load Chat History ```go -// Get conversation list -convs, _ := chatStore.ListConversations(ConversationFilter{ +// Get chat list +chats, _ := chatStore.ListChats(ChatFilter{ UserID: "user123", Status: "active", Page: 1, PageSize: 20, }) -// Get messages for a conversation -messages, _ := chatStore.GetMessages("conv_123", MessageFilter{ +// Get messages for a chat +messages, _ := chatStore.GetMessages("chat_123", MessageFilter{ Limit: 100, }) // Return to frontend return map[string]interface{}{ - "conversation": conv, - "messages": messages, + "chat": chat, + "messages": messages, } ``` @@ -517,18 +586,25 @@ return map[string]interface{}{ ```go func (ast *Assistant) Resume(ctx *Context) error { // 1. Find last incomplete step - step, _ := chatStore.GetLastIncompleteStep(ctx.ConversationID) + step, _ := chatStore.GetLastIncompleteStep(ctx.ChatID) if step == nil { return nil // Nothing to resume } - // 2. Check if this is an A2A nested call + // 2. Restore Space data from snapshot + if step.SpaceSnapshot != nil && ctx.Space != nil { + for key, value := range step.SpaceSnapshot { + ctx.Space.Set(key, value) + } + } + + // 3. Check if this is an A2A nested call if step.StackDepth > 0 { // Need to rebuild the call stack return ast.ResumeNestedCall(ctx, step) } - // 3. Resume based on step type + // 4. Resume based on step type switch step.Type { case "llm": // Re-execute LLM call with saved input @@ -542,6 +618,12 @@ func (ast *Assistant) Resume(ctx *Context) error { case "hook_next": // Re-execute hook return ast.executeHookNext(ctx, step.Input) + + case "delegate": + // Resume delegated agent call + agentID := step.Input["agent_id"].(string) + messages := step.Input["messages"].([]Message) + return ast.delegateToAgent(ctx, agentID, messages) } return nil @@ -590,22 +672,37 @@ When Assistant A delegates to Assistant B, the step records look like: Request: User asks "analyze this data and visualize it" Step Records: -┌─────┬─────────────┬─────────────┬──────────┬────────────┬───────┬─────────────┐ -│ seq │ assistant │ stack_id │ parent │ depth │ type │ status │ -├─────┼─────────────┼─────────────┼──────────┼────────────┼───────┼─────────────┤ -│ 1 │ analyzer │ stk_001 │ null │ 0 │ input │ completed │ -│ 2 │ analyzer │ stk_001 │ null │ 0 │ llm │ completed │ -│ 3 │ analyzer │ stk_001 │ null │ 0 │ delegate │ running │ ← delegating -│ 4 │ visualizer │ stk_002 │ stk_001 │ 1 │ input │ completed │ -│ 5 │ visualizer │ stk_002 │ stk_001 │ 1 │ llm │ interrupted │ ← interrupted here -└─────┴─────────────┴─────────────┴──────────┴────────────┴───────┴─────────────┘ +┌─────┬─────────────┬─────────────┬──────────┬───────┬───────┬─────────────┬─────────────────────────────┐ +│ seq │ assistant │ stack_id │ parent │ depth │ type │ status │ space_snapshot │ +├─────┼─────────────┼─────────────┼──────────┼───────┼───────┼─────────────┼─────────────────────────────┤ +│ 1 │ analyzer │ stk_001 │ null │ 0 │ input │ completed │ {} │ +│ 2 │ analyzer │ stk_001 │ null │ 0 │ llm │ completed │ {} │ +│ 3 │ analyzer │ stk_001 │ null │ 0 │ delegate │ running │ {"choose_prompt": "query"} │ ← Space data set before delegate +│ 4 │ visualizer │ stk_002 │ stk_001 │ 1 │ input │ completed │ {"choose_prompt": "query"} │ +│ 5 │ visualizer │ stk_002 │ stk_001 │ 1 │ llm │ interrupted │ {"choose_prompt": "query"} │ ← interrupted here +└─────┴─────────────┴─────────────┴──────────┴───────┴───────┴─────────────┴─────────────────────────────┘ Resume Flow: 1. Find step with status="interrupted" → step 5 -2. Check stack_depth=1 → nested call -3. Get stack path: [stk_001, stk_002] -4. Resume visualizer assistant with step 5's input -5. When visualizer completes, update step 3 (delegate) to completed +2. Restore Space from space_snapshot: {"choose_prompt": "query"} +3. Check stack_depth=1 → nested call +4. Get stack path: [stk_001, stk_002] +5. Resume visualizer assistant with step 5's input +6. When visualizer completes, update step 3 (delegate) to completed +``` + +**Space Snapshot Use Case (from expense assistant):** + +```typescript +// In Next hook, before delegating to another agent +ctx.space.Set("choose_prompt", "query"); +return { + delegate: { agent_id: "expense", messages: payload.messages }, +}; + +// If interrupted during delegate, Resume will: +// 1. Restore space_snapshot → ctx.space now has "choose_prompt": "query" +// 2. The delegated agent's Create hook can read: ctx.space.GetDel("choose_prompt") ``` ## Migration Notes diff --git a/openapi/request/REQUEST_DESIGN.md b/openapi/request/REQUEST_DESIGN.md index 15b11ea0..cbd20246 100644 --- a/openapi/request/REQUEST_DESIGN.md +++ b/openapi/request/REQUEST_DESIGN.md @@ -284,135 +284,128 @@ type QuotaKey struct { ## Middleware Design -### Request Flow +### Modular Middleware Architecture + +Each middleware is independent and can be composed based on business needs. ``` -Request arrives - │ - ├── 1. Generate request_id (uuid or nanoid) - │ - ├── 2. Set request_id in context and response header - │ c.Set("request_id", requestID) - │ c.Header("X-Request-ID", requestID) - │ - ├── 3. Get auth info from context (set by OAuth Guard) - │ authInfo := authorized.GetInfo(c) - │ - ├── 4. Detect service type from endpoint - │ service := detectService(c.FullPath()) - │ - ├── 5. Create request record (async) - │ status = "running" - │ - ├── 6. Check rate limits - │ if exceeded → return 429, update status = "failed" - │ - ├── 7. Execute handler - │ c.Next() - │ - └── 8. Update request record (async) - status = "completed" or "failed" - duration_ms = time.Since(start) - status_code = c.Writer.Status() +┌─────────────────────────────────────────────────────────────┐ +│ Available Middlewares │ +├─────────────────────────────────────────────────────────────┤ +│ │ +│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ +│ │ RequestID │ │ RateLimit │ │ Quota │ │ +│ │ (Basic) │ │ (Protect) │ │ (Billing) │ │ +│ └─────────────┘ └─────────────┘ └─────────────┘ │ +│ │ +│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ +│ │ Metrics │ │ Archive │ │ Billing │ │ +│ │ (Monitor) │ │ (Audit) │ │ (Charge) │ │ +│ └─────────────┘ └─────────────┘ └─────────────┘ │ +│ │ +└─────────────────────────────────────────────────────────────┘ ``` -### Implementation +### Middleware List + +| Middleware | File | Purpose | Dependencies | +| ----------- | --------------- | -------------------------------- | ------------------ | +| `RequestID` | `request_id.go` | Generate and track request ID | None | +| `RateLimit` | `ratelimit.go` | Request frequency limiting | KV, RequestID | +| `Quota` | `quota.go` | Token quota enforcement | KV, RequestID | +| `Metrics` | `metrics.go` | Request duration, status metrics | RequestID | +| `Archive` | `archive.go` | Persist request to SQL | SQL, RequestID | +| `Billing` | `billing.go` | Token usage tracking & charging | KV, SQL, RequestID | + +### Usage Examples + +#### Example 1: Full Protection (Agent API) ```go +// Agent API needs all protections +agent := api.Group("/chat") +agent.Use( + request.RequestID(), // Generate request_id + request.RateLimit(kv, config), // Rate limiting + request.Quota(kv, config), // Token quota + request.Metrics(), // Duration tracking + request.Archive(sql), // Audit logging + request.Billing(kv, sql), // Token billing +) +agent.POST("/completions", handler.ChatCompletions) +``` + +#### Example 2: Light Protection (File API) + +```go +// File API only needs basic tracking +file := api.Group("/file") +file.Use( + request.RequestID(), // Generate request_id + request.RateLimit(kv, config), // Rate limiting + request.Metrics(), // Duration tracking +) +file.POST("/upload", handler.Upload) +``` + +#### Example 3: Internal API (No Billing) + +```go +// Internal API skips billing +internal := api.Group("/internal") +internal.Use( + request.RequestID(), // Generate request_id + request.Metrics(), // Duration tracking + request.Archive(sql), // Audit logging only +) +internal.GET("/health", handler.Health) +``` + +#### Example 4: Public API (Rate Limit Only) + +```go +// Public endpoints only need rate limiting +public := api.Group("/public") +public.Use( + request.RequestID(), // Generate request_id + request.RateLimit(kv, config), // Rate limiting by IP +) +public.GET("/models", handler.ListModels) +``` + +--- + +### Middleware Implementations + +#### 1. RequestID Middleware (Base) + +```go +// request_id.go package request -import ( - "time" - "github.com/gin-gonic/gin" - "github.com/yaoapp/yao/openapi/oauth/authorized" -) - -// Middleware creates the request tracking middleware -func Middleware(kv KVStore, sql SQLStore) gin.HandlerFunc { +// RequestID generates and sets request ID +func RequestID() gin.HandlerFunc { return func(c *gin.Context) { - startTime := time.Now() - - // 1. Generate request ID requestID := generateRequestID() c.Set("request_id", requestID) c.Header("X-Request-ID", requestID) - // 2. Get auth info - authInfo := authorized.GetInfo(c) + // Also set start time for metrics + c.Set("request_start_time", time.Now()) - // 3. Detect service and resource + // Detect and set service info service := detectService(c.FullPath()) - resourceID := extractResourceID(c, service) + c.Set("request_service", service) + c.Set("request_resource_id", extractResourceID(c, service)) - // 4. KV: Check rate limits (synchronous, must be fast) - if err := checkRateLimit(kv, authInfo, service, c.ClientIP()); err != nil { - c.AbortWithStatusJSON(429, gin.H{ - "error": "rate_limit_exceeded", - "message": err.Error(), - }) - return - } - - // 5. KV: Check quota (synchronous) - if err := checkQuota(kv, authInfo); err != nil { - c.AbortWithStatusJSON(429, gin.H{ - "error": "quota_exceeded", - "message": err.Error(), - }) - return - } - - // 6. KV: Record request status - reqStatus := &RequestStatus{ - RequestID: requestID, - UserID: authInfo.UserID, - TeamID: authInfo.TeamID, - Service: service, - ResourceID: resourceID, - Status: "running", - CreatedAt: startTime, - } - kv.SetRequestStatus(requestID, reqStatus, time.Hour) - - // 7. Execute handler c.Next() - - // 8. KV: Update request status - reqStatus.Status = "completed" - reqStatus.CompletedAt = time.Now() - reqStatus.DurationMs = time.Since(startTime).Milliseconds() - if errMsg := getErrorFromContext(c); errMsg != "" { - reqStatus.Status = "failed" - reqStatus.Error = errMsg - } - kv.SetRequestStatus(requestID, reqStatus, time.Hour) - - // 9. Async: Archive to SQL - go func() { - sql.Archive(&Request{ - RequestID: requestID, - UserID: authInfo.UserID, - TeamID: authInfo.TeamID, - SessionID: authInfo.SessionID, - Endpoint: c.FullPath(), - Method: c.Request.Method, - Service: service, - ResourceID: resourceID, - Status: reqStatus.Status, - StatusCode: c.Writer.Status(), - Referer: c.GetHeader("X-Yao-Referer"), - ClientType: getClientType(c.GetHeader("User-Agent")), - ClientIP: c.ClientIP(), - DurationMs: reqStatus.DurationMs, - Error: reqStatus.Error, - CreatedAt: startTime, - CompletedAt: &reqStatus.CompletedAt, - }) - }() } } -// detectService determines the service type from endpoint +func generateRequestID() string { + return fmt.Sprintf("req_%s", nanoid.New()) +} + func detectService(endpoint string) string { switch { case strings.HasPrefix(endpoint, "/api/chat"): @@ -437,6 +430,256 @@ func detectService(endpoint string) string { } ``` +#### 2. RateLimit Middleware + +```go +// ratelimit.go +package request + +// RateLimit enforces request frequency limits +func RateLimit(kv KVStore, config *RateLimitConfig) gin.HandlerFunc { + return func(c *gin.Context) { + if config == nil || !config.Enabled { + c.Next() + return + } + + authInfo := authorized.GetInfo(c) + service := c.GetString("request_service") + + // Check user rate limit + userKey := fmt.Sprintf("ratelimit:user:%s:%s", authInfo.UserID, service) + userCount, _ := kv.Incr(userKey, 60*time.Second) + if userCount > int64(config.GetUserLimit(service)) { + c.AbortWithStatusJSON(429, gin.H{ + "error": "rate_limit_exceeded", + "message": fmt.Sprintf("User rate limit exceeded: %d requests per minute", config.GetUserLimit(service)), + "retry_after": 60, + }) + return + } + + // Check team rate limit + if authInfo.TeamID != "" { + teamKey := fmt.Sprintf("ratelimit:team:%s:%s", authInfo.TeamID, service) + teamCount, _ := kv.Incr(teamKey, 60*time.Second) + if teamCount > int64(config.GetTeamLimit(service)) { + c.AbortWithStatusJSON(429, gin.H{ + "error": "rate_limit_exceeded", + "message": "Team rate limit exceeded", + "retry_after": 60, + }) + return + } + } + + // Check IP rate limit + ipKey := fmt.Sprintf("ratelimit:ip:%s", c.ClientIP()) + ipCount, _ := kv.Incr(ipKey, 60*time.Second) + if ipCount > int64(config.GetIPLimit()) { + c.AbortWithStatusJSON(429, gin.H{ + "error": "rate_limit_exceeded", + "message": "IP rate limit exceeded", + "retry_after": 60, + }) + return + } + + c.Next() + } +} +``` + +#### 3. Quota Middleware + +```go +// quota.go +package request + +// Quota enforces token quota limits +func Quota(kv KVStore, config *QuotaConfig) gin.HandlerFunc { + return func(c *gin.Context) { + if config == nil || !config.Enabled { + c.Next() + return + } + + authInfo := authorized.GetInfo(c) + + // Check user daily quota + userQuotaKey := fmt.Sprintf("quota:user:%s:daily", authInfo.UserID) + remaining, exists := kv.Get(userQuotaKey) + + if !exists { + // Initialize quota for the day + limit := config.GetUserDailyLimit(authInfo.UserID) + kv.Set(userQuotaKey, limit, 24*time.Hour) + remaining = limit + } + + if remaining <= 0 { + c.AbortWithStatusJSON(429, gin.H{ + "error": "quota_exceeded", + "message": "Daily token quota exceeded", + "reset_at": getNextDayStart(), + }) + return + } + + // Check team monthly quota + if authInfo.TeamID != "" { + teamQuotaKey := fmt.Sprintf("quota:team:%s:monthly", authInfo.TeamID) + teamRemaining, exists := kv.Get(teamQuotaKey) + + if !exists { + limit := config.GetTeamMonthlyLimit(authInfo.TeamID) + kv.Set(teamQuotaKey, limit, 30*24*time.Hour) + teamRemaining = limit + } + + if teamRemaining <= 0 { + c.AbortWithStatusJSON(429, gin.H{ + "error": "quota_exceeded", + "message": "Team monthly token quota exceeded", + "reset_at": getNextMonthStart(), + }) + return + } + } + + c.Next() + } +} +``` + +#### 4. Metrics Middleware + +```go +// metrics.go +package request + +// Metrics tracks request duration and status +func Metrics() gin.HandlerFunc { + return func(c *gin.Context) { + startTime := c.GetTime("request_start_time") + if startTime.IsZero() { + startTime = time.Now() + } + + c.Next() + + // Calculate duration + duration := time.Since(startTime) + c.Set("request_duration_ms", duration.Milliseconds()) + + // Determine status + status := "completed" + if c.Writer.Status() >= 400 { + status = "failed" + } + c.Set("request_status", status) + + // TODO: Export to Prometheus/metrics system + // metrics.RequestDuration.WithLabelValues(service, status).Observe(duration.Seconds()) + // metrics.RequestTotal.WithLabelValues(service, status).Inc() + } +} +``` + +#### 5. Archive Middleware + +```go +// archive.go +package request + +// Archive persists request to SQL for audit +func Archive(sql SQLStore) gin.HandlerFunc { + return func(c *gin.Context) { + c.Next() + + // Get request info from context + requestID := c.GetString("request_id") + if requestID == "" { + return + } + + authInfo := authorized.GetInfo(c) + startTime := c.GetTime("request_start_time") + durationMs := c.GetInt64("request_duration_ms") + status := c.GetString("request_status") + if status == "" { + status = "completed" + } + + completedAt := time.Now() + + // Async archive to SQL + go func() { + sql.Archive(&Request{ + RequestID: requestID, + UserID: authInfo.UserID, + TeamID: authInfo.TeamID, + SessionID: authInfo.SessionID, + Endpoint: c.FullPath(), + Method: c.Request.Method, + Service: c.GetString("request_service"), + ResourceID: c.GetString("request_resource_id"), + Status: status, + StatusCode: c.Writer.Status(), + Referer: c.GetHeader("X-Yao-Referer"), + ClientType: getClientType(c.GetHeader("User-Agent")), + ClientIP: c.ClientIP(), + DurationMs: durationMs, + Error: c.GetString("request_error"), + CreatedAt: startTime, + CompletedAt: &completedAt, + }) + }() + } +} +``` + +#### 6. Billing Middleware + +```go +// billing.go +package request + +// Billing tracks token usage (called by services after completion) +func Billing(kv KVStore, sql SQLStore) gin.HandlerFunc { + return func(c *gin.Context) { + c.Next() + + // Token usage is updated by services via UpdateTokenUsage() + // This middleware just ensures the billing context is available + c.Set("billing_kv", kv) + c.Set("billing_sql", sql) + } +} + +// UpdateTokenUsage is called by services after completion +func UpdateTokenUsage(c *gin.Context, input, output int) error { + kv, ok := c.Get("billing_kv") + if !ok { + return nil // Billing not enabled + } + + sql, _ := c.Get("billing_sql") + requestID := c.GetString("request_id") + authInfo := authorized.GetInfo(c) + + return updateTokenUsageInternal( + kv.(KVStore), + sql.(SQLStore), + requestID, + authInfo.UserID, + authInfo.TeamID, + input, + output, + ) +} +``` + ## Rate Limiting ### Configuration @@ -769,7 +1012,76 @@ type DailyUsage struct { ## Integration with Services -### Agent Service +### Route Registration Example + +```go +// openapi/openapi.go +func (s *OpenAPI) RegisterRoutes(r *gin.Engine) { + api := r.Group("/api") + + // 1. OAuth Guard (authentication) - for all routes + api.Use(oauth.Guard) + + // 2. Register different route groups with different middleware combinations + s.registerAgentRoutes(api) + s.registerKBRoutes(api) + s.registerLLMRoutes(api) + s.registerFileRoutes(api) + s.registerPublicRoutes(api) +} + +func (s *OpenAPI) registerAgentRoutes(api *gin.RouterGroup) { + // Agent API: Full protection + billing + agent := api.Group("/chat") + agent.Use( + request.RequestID(), + request.RateLimit(s.kv, s.rateLimitConfig), + request.Quota(s.kv, s.quotaConfig), + request.Metrics(), + request.Archive(s.sql), + request.Billing(s.kv, s.sql), + ) + agent.POST("/completions", s.handler.ChatCompletions) +} + +func (s *OpenAPI) registerKBRoutes(api *gin.RouterGroup) { + // KB API: Rate limit + archive (no token billing) + kb := api.Group("/kb") + kb.Use( + request.RequestID(), + request.RateLimit(s.kv, s.rateLimitConfig), + request.Metrics(), + request.Archive(s.sql), + ) + kb.POST("/search", s.handler.KBSearch) + kb.POST("/upload", s.handler.KBUpload) +} + +func (s *OpenAPI) registerFileRoutes(api *gin.RouterGroup) { + // File API: Light protection + file := api.Group("/file") + file.Use( + request.RequestID(), + request.RateLimit(s.kv, s.rateLimitConfig), + request.Metrics(), + ) + file.POST("/upload", s.handler.FileUpload) + file.GET("/download/:id", s.handler.FileDownload) +} + +func (s *OpenAPI) registerPublicRoutes(api *gin.RouterGroup) { + // Public API: Rate limit only (no auth required) + public := api.Group("/public") + public.Use( + request.RequestID(), + request.RateLimit(s.kv, s.rateLimitConfig), // IP-based only + ) + public.GET("/models", s.handler.ListModels) + public.GET("/health", s.handler.Health) +} +``` + +### Agent Service Integration ```go // agent/context/openapi.go @@ -780,6 +1092,7 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest // Create context with request ID ctx := New(c.Request.Context(), authInfo, chatID) ctx.RequestID = requestID // Use global request_id + ctx.GinContext = c // Keep gin context for billing // ... } @@ -787,10 +1100,10 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest // agent/assistant/agent.go func (ast *Assistant) Stream(ctx, inputMessages, options) { defer func() { - // Update token usage in global request record - if ctx.RequestID != "" && completionResponse != nil && completionResponse.Usage != nil { + // Update token usage via billing middleware + if ctx.GinContext != nil && completionResponse != nil && completionResponse.Usage != nil { request.UpdateTokenUsage( - ctx.RequestID, + ctx.GinContext, completionResponse.Usage.PromptTokens, completionResponse.Usage.CompletionTokens, ) @@ -801,74 +1114,59 @@ func (ast *Assistant) Stream(ctx, inputMessages, options) { } ``` -### KB Service +### LLM Service Integration ```go -// kb/api/search.go -func (api *API) Search(c *gin.Context) { - requestID := c.GetString("request_id") +// llm/api/completion.go +func (api *API) Completion(c *gin.Context) { + // ... execute LLM call ... - // Perform search... - - // Update metadata if needed - if requestID != "" { - request.UpdateMetadata(requestID, map[string]interface{}{ - "results_count": len(results), - "collection_id": collectionID, - }) + // Update token usage + if response.Usage != nil { + request.UpdateTokenUsage(c, response.Usage.PromptTokens, response.Usage.CompletionTokens) } } ``` -### Middleware Registration - -```go -// openapi/openapi.go -func (s *OpenAPI) RegisterRoutes(r *gin.Engine) { - api := r.Group("/api") - - // 1. OAuth Guard (authentication) - api.Use(oauth.Guard) - - // 2. Request Middleware (tracking, rate limiting) - api.Use(request.Middleware(requestStore)) - - // 3. Service routes - s.registerAgentRoutes(api) - s.registerKBRoutes(api) - s.registerLLMRoutes(api) - // ... -} -``` - ## Summary -### Components +### Middleware Components -| Component | Location | Responsibility | -| ------------ | ------------------------------- | ---------------------------- | -| KV Store | `openapi/request/kv.go` | Real-time: rate limit, quota | -| SQL Store | `openapi/request/sql.go` | Archive: billing, audit | -| Middleware | `openapi/request/middleware.go` | Track requests, orchestrate | -| Rate Limiter | `openapi/request/ratelimit.go` | Enforce rate limits | -| Types | `openapi/request/types.go` | Data structures | +| Middleware | File | Purpose | Storage | +| ----------- | --------------- | ------------------------ | -------- | +| `RequestID` | `request_id.go` | Generate request ID | - | +| `RateLimit` | `ratelimit.go` | Frequency limiting | KV | +| `Quota` | `quota.go` | Token quota enforcement | KV | +| `Metrics` | `metrics.go` | Duration/status tracking | - | +| `Archive` | `archive.go` | Persist to SQL | SQL | +| `Billing` | `billing.go` | Token usage tracking | KV + SQL | -### Storage Comparison +### Storage Components -| Operation | KV (Redis) | SQL (Archive) | -| ---------------- | -------------- | ------------- | -| Rate limit check | ✅ Synchronous | ❌ Not used | -| Quota check | ✅ Synchronous | ❌ Not used | -| Request status | ✅ Synchronous | ❌ Not used | -| Token update | ✅ Synchronous | ✅ Async | -| Billing report | ❌ Not used | ✅ Query | -| Audit log | ❌ Not used | ✅ Query | +| Component | File | Purpose | +| --------- | ---------- | ---------------------------- | +| KV Store | `kv.go` | Real-time: rate limit, quota | +| SQL Store | `sql.go` | Archive: billing, audit | +| Types | `types.go` | Data structures | + +### Middleware Combinations by Use Case + +| Use Case | RequestID | RateLimit | Quota | Metrics | Archive | Billing | +| ------------ | --------- | --------- | ----- | ------- | ------- | ------- | +| Agent API | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| LLM API | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| KB API | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | +| File API | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | +| Public API | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | +| Internal API | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ### Key Points -1. **Two-layer storage**: KV for real-time, SQL for archive -2. **KV operations are synchronous**: Rate limit and quota checks must be fast -3. **SQL writes are async**: Archive happens in background goroutine -4. **Services update tokens via `request_id`**: Updates both KV and SQL -5. **KV data has TTL**: Auto-expires to prevent memory bloat -6. **SQL data is permanent**: For billing and compliance +1. **Modular design**: Each middleware is independent and composable +2. **Business-driven composition**: Routes choose which middleware to use +3. **Two-layer storage**: KV for real-time, SQL for archive +4. **KV operations are synchronous**: Rate limit and quota checks must be fast +5. **SQL writes are async**: Archive happens in background goroutine +6. **Services update tokens via gin context**: `request.UpdateTokenUsage(c, input, output)` +7. **KV data has TTL**: Auto-expires to prevent memory bloat +8. **SQL data is permanent**: For billing and compliance