feat: add Responses API compatibility

2026-07-10 23:33:15 +08:00
parent b7351f3b9b
commit 8847d973a8
24 changed files with 5922 additions and 169 deletions
@@ -31,7 +31,18 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType
       COALESCE(p.dynamic_priority, p.priority) AS effective_priority,
       m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''),
       COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), m.model_name, COALESCE(m.model_alias, ''),
-       $2::text AS requested_model_type, m.display_name, m.capabilities, m.capability_override,
+       $2::text AS requested_model_type, m.display_name,
+       CASE
+         WHEN b.capabilities #> '{text_generate,supportedApiProtocols}' IS NOT NULL
+         THEN jsonb_set(
+           COALESCE(m.capabilities, '{}'::jsonb),
+           '{text_generate,supportedApiProtocols}',
+           b.capabilities #> '{text_generate,supportedApiProtocols}',
+           true
+         )
+         ELSE m.capabilities
+       END AS effective_capabilities,
+       m.capability_override,
       COALESCE(b.base_billing_config, '{}'::jsonb), m.billing_config, m.billing_config_override,
       m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''),
       COALESCE(b.pricing_rule_set_id::text, ''),
@@ -0,0 +1,183 @@
+package store
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"strings"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+)
+
+var ErrInvalidPreviousResponseID = errors.New("invalid previous response id")
+var ErrResponseChainTooDeep = errors.New("response chain exceeds the maximum depth")
+
+const maxResponseChainHistory = 100
+
+type ResponseChain struct {
+	PublicResponseID   string
+	ParentResponseID   string
+	GatewayTaskID      string
+	RequestedModel     string
+	GatewayUserID      string
+	UserID             string
+	UserSource         string
+	GatewayTenantID    string
+	TenantID           string
+	TenantKey          string
+	APIKeyID           string
+	PlatformID         string
+	PlatformModelID    string
+	ClientID           string
+	UpstreamProtocol   string
+	UpstreamEndpoint   string
+	UpstreamResponseID string
+	RequestSnapshot    map[string]any
+	ResponseSnapshot   map[string]any
+	InternalSnapshot   map[string]any
+	CreatedAt          time.Time
+	ExpiresAt          time.Time
+}
+
+type CreateResponseChainInput struct {
+	ResponseChain
+}
+
+func (s *Store) CreateResponseChain(ctx context.Context, input CreateResponseChainInput) error {
+	requestJSON, err := json.Marshal(input.RequestSnapshot)
+	if err != nil {
+		return err
+	}
+	responseJSON, err := json.Marshal(input.ResponseSnapshot)
+	if err != nil {
+		return err
+	}
+	internalJSON, err := json.Marshal(input.InternalSnapshot)
+	if err != nil {
+		return err
+	}
+	_, err = s.pool.Exec(ctx, `
+INSERT INTO gateway_response_chains (
+  public_response_id, parent_response_id, gateway_task_id,
+  gateway_user_id, user_id, user_source, gateway_tenant_id, tenant_id, tenant_key, api_key_id,
+  platform_id, platform_model_id, client_id, upstream_protocol, upstream_endpoint, upstream_response_id,
+  request_snapshot, response_snapshot, internal_snapshot, expires_at
+) VALUES (
+  $1, NULLIF($2, ''), $3::uuid,
+  NULLIF($4, '')::uuid, $5, $6, NULLIF($7, '')::uuid, NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''),
+  $11::uuid, $12::uuid, $13, $14, $15, NULLIF($16, ''),
+  $17::jsonb, $18::jsonb, $19::jsonb, $20
+)
+ON CONFLICT (public_response_id) DO NOTHING`,
+		input.PublicResponseID, input.ParentResponseID, input.GatewayTaskID,
+		input.GatewayUserID, input.UserID, input.UserSource, input.GatewayTenantID, input.TenantID, input.TenantKey, input.APIKeyID,
+		input.PlatformID, input.PlatformModelID, input.ClientID, input.UpstreamProtocol, input.UpstreamEndpoint, input.UpstreamResponseID,
+		requestJSON, responseJSON, internalJSON, input.ExpiresAt,
+	)
+	return err
+}
+
+func (s *Store) GetResponseChain(ctx context.Context, publicResponseID string, owner GatewayTask) (ResponseChain, error) {
+	chain, err := s.scanResponseChain(s.pool.QueryRow(ctx, `
+	SELECT chain.public_response_id, COALESCE(chain.parent_response_id, ''), chain.gateway_task_id::text,
+	       task.requested_model,
+	       COALESCE(chain.gateway_user_id::text, ''), chain.user_id, chain.user_source,
+	       COALESCE(chain.gateway_tenant_id::text, ''), COALESCE(chain.tenant_id, ''), COALESCE(chain.tenant_key, ''), COALESCE(chain.api_key_id, ''),
+	       chain.platform_id::text, chain.platform_model_id::text, chain.client_id, chain.upstream_protocol, chain.upstream_endpoint,
+	       COALESCE(chain.upstream_response_id, ''), chain.request_snapshot, chain.response_snapshot, chain.internal_snapshot,
+	       chain.created_at, chain.expires_at
+FROM gateway_response_chains chain
+JOIN gateway_tasks task ON task.id = chain.gateway_task_id
+WHERE chain.public_response_id = $1`, strings.TrimSpace(publicResponseID)))
+	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return ResponseChain{}, ErrInvalidPreviousResponseID
+		}
+		return ResponseChain{}, err
+	}
+	if chain.ExpiresAt.Before(time.Now()) || !sameResponseChainOwner(chain, owner) {
+		return ResponseChain{}, ErrInvalidPreviousResponseID
+	}
+	return chain, nil
+}
+
+func (s *Store) ListResponseChainHistory(ctx context.Context, current ResponseChain) ([]ResponseChain, error) {
+	rows, err := s.pool.Query(ctx, `
+WITH RECURSIVE history AS (
+  SELECT chain.*, 0 AS depth
+  FROM gateway_response_chains chain
+  WHERE chain.public_response_id = $1
+  UNION ALL
+  SELECT parent.*, history.depth + 1
+  FROM gateway_response_chains parent
+  JOIN history ON parent.public_response_id = history.parent_response_id
+	  WHERE history.depth < $2
+)
+	SELECT history.public_response_id, COALESCE(history.parent_response_id, ''), history.gateway_task_id::text,
+	       task.requested_model,
+	       COALESCE(history.gateway_user_id::text, ''), history.user_id, history.user_source,
+	       COALESCE(history.gateway_tenant_id::text, ''), COALESCE(history.tenant_id, ''), COALESCE(history.tenant_key, ''), COALESCE(history.api_key_id, ''),
+	       history.platform_id::text, history.platform_model_id::text, history.client_id, history.upstream_protocol, history.upstream_endpoint,
+	       COALESCE(history.upstream_response_id, ''), history.request_snapshot, history.response_snapshot, history.internal_snapshot,
+	       history.created_at, history.expires_at
+FROM history
+JOIN gateway_tasks task ON task.id = history.gateway_task_id
+ORDER BY depth DESC`, current.PublicResponseID, maxResponseChainHistory)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	out := make([]ResponseChain, 0)
+	for rows.Next() {
+		item, scanErr := s.scanResponseChain(rows)
+		if scanErr != nil {
+			return nil, scanErr
+		}
+		out = append(out, item)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	if len(out) > maxResponseChainHistory {
+		return nil, ErrResponseChainTooDeep
+	}
+	return out, nil
+}
+
+type responseChainScanner interface {
+	Scan(dest ...any) error
+}
+
+func (s *Store) scanResponseChain(scanner responseChainScanner) (ResponseChain, error) {
+	var item ResponseChain
+	var requestJSON, responseJSON, internalJSON []byte
+	err := scanner.Scan(
+		&item.PublicResponseID, &item.ParentResponseID, &item.GatewayTaskID, &item.RequestedModel,
+		&item.GatewayUserID, &item.UserID, &item.UserSource,
+		&item.GatewayTenantID, &item.TenantID, &item.TenantKey, &item.APIKeyID,
+		&item.PlatformID, &item.PlatformModelID, &item.ClientID,
+		&item.UpstreamProtocol, &item.UpstreamEndpoint, &item.UpstreamResponseID,
+		&requestJSON, &responseJSON, &internalJSON, &item.CreatedAt, &item.ExpiresAt,
+	)
+	if err != nil {
+		return ResponseChain{}, err
+	}
+	item.RequestSnapshot = decodeObject(requestJSON)
+	item.ResponseSnapshot = decodeObject(responseJSON)
+	item.InternalSnapshot = decodeObject(internalJSON)
+	return item, nil
+}
+
+func sameResponseChainOwner(chain ResponseChain, task GatewayTask) bool {
+	if chain.GatewayUserID != "" || task.GatewayUserID != "" {
+		if chain.GatewayUserID == "" || chain.GatewayUserID != task.GatewayUserID {
+			return false
+		}
+	} else if chain.UserID != task.UserID || chain.UserSource != task.UserSource {
+		return false
+	}
+	return chain.GatewayTenantID == task.GatewayTenantID &&
+		chain.TenantID == task.TenantID &&
+		chain.TenantKey == task.TenantKey
+}
@@ -0,0 +1,36 @@
+package store
+
+import "testing"
+
+func TestSameResponseChainOwnerIsolatesUserAndTenant(t *testing.T) {
+	chain := ResponseChain{
+		GatewayUserID: "user-1", GatewayTenantID: "tenant-1", TenantID: "external-tenant", TenantKey: "tenant-key",
+	}
+	owner := GatewayTask{
+		GatewayUserID: "user-1", GatewayTenantID: "tenant-1", TenantID: "external-tenant", TenantKey: "tenant-key",
+	}
+	if !sameResponseChainOwner(chain, owner) {
+		t.Fatal("matching user and tenant should own the response chain")
+	}
+	owner.GatewayUserID = "user-2"
+	if sameResponseChainOwner(chain, owner) {
+		t.Fatal("another user must not access the response chain")
+	}
+	owner.GatewayUserID = "user-1"
+	owner.GatewayTenantID = "tenant-2"
+	if sameResponseChainOwner(chain, owner) {
+		t.Fatal("another tenant must not access the response chain")
+	}
+}
+
+func TestSameResponseChainOwnerSupportsExternalIdentityFallback(t *testing.T) {
+	chain := ResponseChain{UserID: "external-user", UserSource: "server-main", TenantID: "tenant-a"}
+	owner := GatewayTask{UserID: "external-user", UserSource: "server-main", TenantID: "tenant-a"}
+	if !sameResponseChainOwner(chain, owner) {
+		t.Fatal("matching external identity should own the response chain")
+	}
+	owner.UserSource = "gateway"
+	if sameResponseChainOwner(chain, owner) {
+		t.Fatal("identity source must participate in isolation")
+	}
+}
@@ -167,6 +167,7 @@ type RuntimeModelCandidate struct {
 	RunningCount             float64
 	WaitingCount             float64
 	LastAssignedUnix         float64
+	ResponseProtocol         string
 }

 type RuntimeCandidateCacheAffinity struct {