mirror of
https://github.com/CJackHwang/ds2api.git
synced 2026-05-10 19:27:41 +08:00
fix(vercel): enable auto-delete session on Vercel stream release
The "delete current conversation" feature was not working on Vercel deployment because the stream flow uses a separate lease mechanism. The session_id created during prepare phase was not preserved for deletion when the stream ends. Changes: - Add SessionID field to streamLease struct to preserve session_id - Pass session_id to holdStreamLease during prepare - Modify releaseStreamLease to return auth and session_id - Call autoDeleteRemoteSession in handleVercelStreamRelease when releasing a lease with auto-delete mode enabled Closes #vercel-auto-delete
This commit is contained in:
@@ -33,6 +33,7 @@ type Handler struct {
|
||||
|
||||
type streamLease struct {
|
||||
Auth *auth.RequestAuth
|
||||
SessionID string
|
||||
ExpiresAt time.Time
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -64,14 +65,14 @@ func TestVercelInternalSecret(t *testing.T) {
|
||||
|
||||
func TestStreamLeaseLifecycle(t *testing.T) {
|
||||
h := &Handler{}
|
||||
leaseID := h.holdStreamLease(&auth.RequestAuth{UseConfigToken: false})
|
||||
leaseID := h.holdStreamLease(&auth.RequestAuth{UseConfigToken: false}, "test-session-id")
|
||||
if leaseID == "" {
|
||||
t.Fatalf("expected non-empty lease id")
|
||||
}
|
||||
if ok := h.releaseStreamLease(leaseID); !ok {
|
||||
if ok, _, _ := h.releaseStreamLease(leaseID); !ok {
|
||||
t.Fatalf("expected lease release success")
|
||||
}
|
||||
if ok := h.releaseStreamLease(leaseID); ok {
|
||||
if ok, _, _ := h.releaseStreamLease(leaseID); ok {
|
||||
t.Fatalf("expected duplicate release to fail")
|
||||
}
|
||||
}
|
||||
@@ -141,6 +142,94 @@ func TestHandleVercelStreamPrepareAppliesCurrentInputFile(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
type vercelReleaseAutoDeleteDSStub struct {
|
||||
resp *http.Response
|
||||
deleteCallCount int
|
||||
deletedSessionID string
|
||||
deletedToken string
|
||||
deleteErr error
|
||||
}
|
||||
|
||||
func (m *vercelReleaseAutoDeleteDSStub) CreateSession(_ context.Context, _ *auth.RequestAuth, _ int) (string, error) {
|
||||
return "session-id", nil
|
||||
}
|
||||
|
||||
func (m *vercelReleaseAutoDeleteDSStub) GetPow(_ context.Context, _ *auth.RequestAuth, _ int) (string, error) {
|
||||
return "pow", nil
|
||||
}
|
||||
|
||||
func (m *vercelReleaseAutoDeleteDSStub) UploadFile(_ context.Context, _ *auth.RequestAuth, _ dsclient.UploadFileRequest, _ int) (*dsclient.UploadFileResult, error) {
|
||||
return &dsclient.UploadFileResult{ID: "file-id", Filename: "file.txt", Bytes: 1, Status: "uploaded"}, nil
|
||||
}
|
||||
|
||||
func (m *vercelReleaseAutoDeleteDSStub) CallCompletion(_ context.Context, _ *auth.RequestAuth, _ map[string]any, _ string, _ int) (*http.Response, error) {
|
||||
return m.resp, nil
|
||||
}
|
||||
|
||||
func (m *vercelReleaseAutoDeleteDSStub) DeleteSessionForToken(_ context.Context, token string, sessionID string) (*dsclient.DeleteSessionResult, error) {
|
||||
m.deleteCallCount++
|
||||
m.deletedSessionID = sessionID
|
||||
m.deletedToken = token
|
||||
if m.deleteErr != nil {
|
||||
return nil, m.deleteErr
|
||||
}
|
||||
return &dsclient.DeleteSessionResult{SessionID: sessionID, Success: true}, nil
|
||||
}
|
||||
|
||||
func (m *vercelReleaseAutoDeleteDSStub) DeleteAllSessionsForToken(_ context.Context, _ string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type vercelReleaseAuthStub struct{}
|
||||
|
||||
func (a *vercelReleaseAuthStub) Determine(_ *http.Request) (*auth.RequestAuth, error) {
|
||||
return &auth.RequestAuth{DeepSeekToken: "test-token", AccountID: "test-account"}, nil
|
||||
}
|
||||
|
||||
func (a *vercelReleaseAuthStub) DetermineCaller(_ *http.Request) (*auth.RequestAuth, error) {
|
||||
return &auth.RequestAuth{DeepSeekToken: "test-token", AccountID: "test-account"}, nil
|
||||
}
|
||||
|
||||
func (a *vercelReleaseAuthStub) Release(_ *auth.RequestAuth) {}
|
||||
|
||||
func TestHandleVercelStreamReleaseTriggersAutoDelete(t *testing.T) {
|
||||
t.Setenv("VERCEL", "1")
|
||||
t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "stream-secret")
|
||||
|
||||
ds := &vercelReleaseAutoDeleteDSStub{}
|
||||
h := &Handler{
|
||||
Store: mockOpenAIConfig{
|
||||
autoDeleteMode: "single",
|
||||
},
|
||||
Auth: &vercelReleaseAuthStub{},
|
||||
DS: ds,
|
||||
}
|
||||
|
||||
leaseID := h.holdStreamLease(&auth.RequestAuth{DeepSeekToken: "test-token", AccountID: "test-account"}, "session-to-delete")
|
||||
if leaseID == "" {
|
||||
t.Fatalf("expected non-empty lease id")
|
||||
}
|
||||
|
||||
reqBody := map[string]any{"lease_id": leaseID}
|
||||
reqJSON, _ := json.Marshal(reqBody)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions?__stream_release=1", strings.NewReader(string(reqJSON)))
|
||||
req.Header.Set("X-Ds2-Internal-Token", "stream-secret")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
h.handleVercelStreamRelease(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if ds.deleteCallCount != 1 {
|
||||
t.Fatalf("expected auto delete call count=1, got %d", ds.deleteCallCount)
|
||||
}
|
||||
if ds.deletedSessionID != "session-to-delete" {
|
||||
t.Fatalf("expected deleted session id=session-to-delete, got %q", ds.deletedSessionID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleVercelStreamPrepareMapsCurrentInputFileManagedAuthFailureTo401(t *testing.T) {
|
||||
t.Setenv("VERCEL", "1")
|
||||
t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "stream-secret")
|
||||
|
||||
@@ -96,7 +96,7 @@ func (h *Handler) handleVercelStreamPrepare(w http.ResponseWriter, r *http.Reque
|
||||
}
|
||||
|
||||
payload := stdReq.CompletionPayload(sessionID)
|
||||
leaseID := h.holdStreamLease(a)
|
||||
leaseID := h.holdStreamLease(a, sessionID)
|
||||
if leaseID == "" {
|
||||
writeOpenAIError(w, http.StatusInternalServerError, "failed to create stream lease")
|
||||
return
|
||||
@@ -140,10 +140,14 @@ func (h *Handler) handleVercelStreamRelease(w http.ResponseWriter, r *http.Reque
|
||||
writeOpenAIError(w, http.StatusBadRequest, "lease_id is required")
|
||||
return
|
||||
}
|
||||
if !h.releaseStreamLease(leaseID) {
|
||||
ok, leaseAuth, sessionID := h.releaseStreamLease(leaseID)
|
||||
if !ok {
|
||||
writeOpenAIError(w, http.StatusNotFound, "stream lease not found")
|
||||
return
|
||||
}
|
||||
if sessionID != "" && leaseAuth != nil {
|
||||
h.autoDeleteRemoteSession(r.Context(), leaseAuth, sessionID)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"success": true})
|
||||
}
|
||||
|
||||
@@ -216,7 +220,7 @@ func vercelInternalSecret() string {
|
||||
return "admin"
|
||||
}
|
||||
|
||||
func (h *Handler) holdStreamLease(a *auth.RequestAuth) string {
|
||||
func (h *Handler) holdStreamLease(a *auth.RequestAuth, sessionID string) string {
|
||||
if a == nil {
|
||||
return ""
|
||||
}
|
||||
@@ -234,6 +238,7 @@ func (h *Handler) holdStreamLease(a *auth.RequestAuth) string {
|
||||
leaseID := newLeaseID()
|
||||
h.streamLeases[leaseID] = streamLease{
|
||||
Auth: a,
|
||||
SessionID: sessionID,
|
||||
ExpiresAt: now.Add(ttl),
|
||||
}
|
||||
h.leaseMu.Unlock()
|
||||
@@ -255,10 +260,10 @@ func (h *Handler) lookupStreamLeaseAuth(leaseID string) *auth.RequestAuth {
|
||||
return lease.Auth
|
||||
}
|
||||
|
||||
func (h *Handler) releaseStreamLease(leaseID string) bool {
|
||||
func (h *Handler) releaseStreamLease(leaseID string) (bool, *auth.RequestAuth, string) {
|
||||
leaseID = strings.TrimSpace(leaseID)
|
||||
if leaseID == "" {
|
||||
return false
|
||||
return false, nil, ""
|
||||
}
|
||||
|
||||
h.leaseMu.Lock()
|
||||
@@ -271,12 +276,12 @@ func (h *Handler) releaseStreamLease(leaseID string) bool {
|
||||
h.releaseExpiredAuths(expired)
|
||||
|
||||
if !ok {
|
||||
return false
|
||||
return false, nil, ""
|
||||
}
|
||||
if h.Auth != nil {
|
||||
h.Auth.Release(lease.Auth)
|
||||
}
|
||||
return true
|
||||
return true, lease.Auth, lease.SessionID
|
||||
}
|
||||
|
||||
func (h *Handler) popExpiredLeasesLocked(now time.Time) []*auth.RequestAuth {
|
||||
|
||||
Reference in New Issue
Block a user