fix(btw): 修复长对话上下文预算与旁路输入布局

This commit is contained in:
ji ruo committed 2026-09-11 00:41:01 +08:00
1 parent 555539bc8e
commit d893861d00
17 files changed
+1295 -92

No files matched your search

+35 -27
View File
@@ -9,6 +9,7 @@ import (
"fmt"
"log"
"net/http"
"os"
"strconv"
"strings"
"sync"
@@ -29,13 +30,14 @@ type sideRun struct {
done chan struct{}
}
type sideQuestionState struct {
mu sync.Mutex
commands sync.Mutex // short admission / clear operations only, never inference
pending map[string]sidequestion.Snapshot
latest map[string]sidequestion.Snapshot
seen map[string][2]int64
runs map[string]sideRun
done chan struct{}
mu sync.Mutex
commands sync.Mutex // short admission / clear operations only, never inference
pending map[string]sidequestion.Snapshot
latest map[string]sidequestion.Snapshot
seen map[string][2]int64
runs map[string]sideRun
done chan struct{}
outputTokens int
}
func sideModel(cfg agent.Config, id int64, name string) sidequestion.Model {
@@ -50,6 +52,9 @@ func bindSideProvider(p llm.Provider, cfg agent.Config, id int64, name string) l
func (s *Server) initSideQuestions() {
s.side = &sideQuestionState{pending: map[string]sidequestion.Snapshot{}, latest: map[string]sidequestion.Snapshot{}, seen: map[string][2]int64{}, runs: map[string]sideRun{}, done: make(chan struct{})}
if n, err := strconv.Atoi(os.Getenv("ARTEX_BTW_MAX_OUTPUT_TOKENS")); err == nil && n >= 256 && n <= 32768 {
s.side.outputTokens = n
}
if s.m.pg == nil {
close(s.side.done)
return
@@ -394,16 +399,6 @@ func (s *Server) handleSideQuestions(w http.ResponseWriter, r *http.Request, p s
writeErr(w, 500, err.Error())
return
}
history, err := s.m.pg.SideReplay(r.Context(), key)
if err != nil {
writeErr(w, 500, err.Error())
return
}
req, err := sidequestion.BuildRequest(*snap, history, in.Question)
if err != nil {
writeErr(w, 409, err.Error())
return
}
e, created, err := s.m.pg.StartSideRequest(r.Context(), *snap, in.ClientID, in.Question)
if err != nil {
writeErr(w, 409, err.Error())
@@ -414,12 +409,12 @@ func (s *Server) handleSideQuestions(w http.ResponseWriter, r *http.Request, p s
s.side.mu.Lock()
s.side.runs[e.ID] = sideRun{key: key, parent: p, cancel: cancel, done: make(chan struct{})}
s.side.mu.Unlock()
go s.runSide(ctx, cancel, *e, p, provider, req)
go s.runSide(ctx, cancel, *e, p, provider, *snap)
}
writeJSON(w, 202, e)
}
func (s *Server) runSide(ctx context.Context, cancel context.CancelFunc, e sidequestion.Exchange, parent sidequestion.Parent, provider llm.Provider, req llm.CompletionRequest) {
func (s *Server) runSide(ctx context.Context, cancel context.CancelFunc, e sidequestion.Exchange, parent sidequestion.Parent, provider llm.Provider, snapshot sidequestion.Snapshot) {
defer cancel()
defer func() {
s.side.mu.Lock()
@@ -463,15 +458,28 @@ func (s *Server) runSide(ctx context.Context, cancel context.CancelFunc, e sideq
}
return ok, nil
}
answer, runErr := (sidequestion.SideQuestionService{Provider: provider}).Answer(ctx, req, e.Model.Streaming, func(a sidequestion.Answer) {
e.Answer, e.Usage = a.Text, a.Usage
if time.Since(lastSave) >= 250*time.Millisecond {
lastSave = time.Now()
if ok, _ := persist(); !ok {
cancel()
memory, runErr := s.m.pg.SideMemory(ctx, e)
var answer sidequestion.Answer
if runErr == nil {
answer, e.Context, runErr = (sidequestion.SideQuestionService{Provider: provider}).Respond(ctx, snapshot, e.Question, sidequestion.Replay{
Memory: memory,
Load: func(ctx context.Context, after int64) ([]sidequestion.Exchange, error) {
return s.m.pg.SideReplayPage(ctx, e, after)
},
Save: func(ctx context.Context, memory sidequestion.Memory) error {
return s.m.pg.SaveSideMemory(ctx, e, memory)
},
}, sidequestion.ContextOptions{OutputTokens: s.side.outputTokens}, func(a sidequestion.Answer, info sidequestion.ContextInfo) {
phaseChanged := e.Context.Phase != info.Phase
e.Answer, e.Usage, e.Context = a.Text, a.Usage, info
if phaseChanged || time.Since(lastSave) >= 250*time.Millisecond {
lastSave = time.Now()
if ok, _ := persist(); !ok {
cancel()
}
}
}
})
})
}
e.Answer, e.Usage = answer.Text, answer.Usage
e.Status = "completed"
if ctx.Err() != nil {
+80 -2
View File
@@ -24,6 +24,7 @@ import (
type sideHTTPProvider struct {
started chan llm.CompletionRequest
release chan struct{}
summary func(context.Context, llm.CompletionRequest) (llm.Message, llm.Usage, error)
}
func (p *sideHTTPProvider) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
@@ -44,6 +45,10 @@ func (p *sideHTTPProvider) Stream(ctx context.Context, req llm.CompletionRequest
}
}
func (p *sideHTTPProvider) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
if req.Thinking == "disabled" && p.summary != nil {
msg, usage, err := p.summary(ctx, req)
return msg, "end_turn", usage, err
}
for _, err := range p.Stream(ctx, req) {
if err != nil {
return llm.Message{}, "", llm.Usage{}, err
@@ -52,6 +57,79 @@ func (p *sideHTTPProvider) Complete(ctx context.Context, req llm.CompletionReque
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("atomic answer")}}, "end_turn", llm.Usage{InputTokens: 19}, nil
}
func TestSideHTTPPreparationCancellationAndClear(t *testing.T) {
for _, clear := range []bool{false, true} {
t.Run(fmt.Sprint(clear), func(t *testing.T) {
f := newSideHTTPFixture(t)
p, path := f.conversation(t)
snap := f.checkpoint(t, p)
for i := 0; i < 21; i++ {
e, _, err := f.m.pg.StartSideRequest(t.Context(), snap, fmt.Sprint(i), "past question")
if err != nil {
t.Fatal(err)
}
e.Answer = "old answer"
e.Sequence = 1
e.Status = "completed"
if _, err = f.m.pg.UpdateSideRequest(t.Context(), *e); err != nil {
t.Fatal(err)
}
}
started := make(chan struct{})
f.provider.summary = func(ctx context.Context, req llm.CompletionRequest) (llm.Message, llm.Usage, error) {
if len(req.Tools) != 0 {
t.Error("summary has tools")
}
close(started)
<-ctx.Done()
return llm.UserText("late summary"), llm.Usage{InputTokens: 17}, nil
}
e := decodeSide(t, f.call(t, "POST", path, map[string]string{"question": "follow-up", "client_request_id": "new"}, 202))
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("summary did not start")
}
row, err := f.m.pg.SideRequest(t.Context(), e.ID)
if err != nil || row.Context.Phase != "summarizing_history" {
t.Fatalf("phase not persisted: %+v %v", row, err)
}
// A blocked summary must not hold the global admission/clear mutex.
other, otherPath := f.conversation(t)
f.checkpoint(t, other)
f.call(t, "POST", otherPath, map[string]string{"question": "other session", "client_request_id": "independent"}, 202)
f.s.side.mu.Lock()
done := f.s.side.runs[e.ID].done
f.s.side.mu.Unlock()
if clear {
f.call(t, "DELETE", path, nil, 200)
} else {
f.call(t, "POST", "/api/side-questions/"+e.ID+"/cancel", nil, 200)
}
select {
case <-done:
case <-time.After(time.Second):
t.Fatal("summary cancellation did not settle")
}
row, err = f.m.pg.SideRequest(t.Context(), e.ID)
if err != nil {
t.Fatal(err)
}
if clear {
if row != nil {
t.Fatal("cleared request reappeared")
}
} else if row.Status != "cancelled" || row.Usage.InputTokens != 17 {
t.Fatalf("cancel usage: %+v", row)
}
var memory string
if err = f.m.pg.QueryRow(`SELECT memory::text FROM side_question_sessions WHERE session_key=$1`, p.Key()).Scan(&memory); err != nil || memory != "{}" {
t.Fatalf("late summary saved: %s %v", memory, err)
}
})
}
}
type sideHTTPFixture struct {
s *Server
m *Manager
@@ -225,11 +303,11 @@ func TestSideHTTPBusyIsolationClearAndReconnect(t *testing.T) {
func TestSideHTTPGlobalLimitTaskWorkerAndDeletion(t *testing.T) {
f := newSideHTTPFixture(t)
var requests []sidequestion.Exchange
for i := 0; i < 5; i++ {
for i := 0; i < 20; i++ {
p, path := f.conversation(t)
f.checkpoint(t, p)
want := 202
if i == 4 {
if i >= 4 {
want = 429
}
w := f.call(t, "POST", path, map[string]string{"question": "blocked", "client_request_id": "id"}, want)