opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor chat storage design and enhance data models

Browse source 
- 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.
This commit is contained in:
Max 2025-12-08 18:47:45 +08:00
parent 467d4e2398
commit 42d13ec1cb
2 changed files with 716 additions and 321 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

235
agent/store/CHAT_STORAGE_DESIGN.md
View file

@ -40,11 +40,11 @@ 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 sessions | `agent/store` | `agent_chat` |
| Chat messages | `agent/store` | `agent_message` |
| Execution steps | `agent/store` | `agent_step` |
@ -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 |
| `chat_id` | string(64) | No | Unique | Unique chat identifier |
| `title` | string(500) | Yes | - | Chat 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") |
| `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 |
| -------------------- | ----------------- | ----- |
| `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
@ -115,10 +140,10 @@ 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 |
| `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.) |
@ -126,7 +151,7 @@ Stores user-visible messages (both user input and assistant responses).
| `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 |
| `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 |
@ -134,8 +159,8 @@ Stores user-visible messages (both user input and assistant responses).
**Indexes:**
| Name | Columns | Type |
| ------------------- | ----------------------------- | ----- |
| `idx_msg_conv_seq` | `conversation_id`, `sequence` | index |
| ------------------- | --------------------- | ----- |
| `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 |
@ -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,9 +397,15 @@ 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
ChatID: ctx.ChatID, // ChatID
RequestID: ctx.RequestID, // From OpenAPI middleware
AssistantID: ctx.AssistantID,
StackID: ctx.Stack.ID,
@ -362,6 +415,7 @@ func createStep(ctx *Context, stepType, status string, input, output interface{}
Status: status,
Input: input,
Output: output,
SpaceSnapshot: spaceSnapshot, // Shared space data for recovery
Sequence: nextSequence(),
}
}
@ -373,39 +427,53 @@ 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"`
// Chat represents a chat session
type Chat struct {
ChatID string `json:"chat_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"`
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"`
@ -415,7 +483,7 @@ type Conversation struct {
// Message represents a chat message
type Message struct {
MessageID string `json:"message_id"`
ConversationID string `json:"conversation_id"`
ChatID string `json:"chat_id"`
RequestID string `json:"request_id,omitempty"`
Role string `json:"role"`
Type string `json:"type"`
@ -432,7 +500,7 @@ type Message struct {
// Step represents an execution step
type Step struct {
StepID string `json:"step_id"`
ConversationID string `json:"conversation_id"`
ChatID string `json:"chat_id"`
RequestID string `json:"request_id"`
AssistantID string `json:"assistant_id"`
StackID string `json:"stack_id"`
@ -442,6 +510,7 @@ type Step struct {
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"`
@ -453,8 +522,8 @@ type Step struct {
### 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,9 +542,9 @@ type MessageFilter struct {
Offset int `json:"offset,omitempty"`
}
// ConversationList paginated response
type ConversationList struct {
Data []*Conversation `json:"data"`
// ChatList paginated response
type ChatList struct {
Data []*Chat `json:"data"`
Page int `json:"page"`
PageSize int `json:"pagesize"`
PageCount int `json:"pagecount"`
@ -492,22 +561,22 @@ 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,
"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

636
openapi/request/REQUEST_DESIGN.md
View file

@ -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