diff --git a/internal/agent/core/process.go b/internal/agent/core/process.go index 529a036..1649d16 100644 --- a/internal/agent/core/process.go +++ b/internal/agent/core/process.go @@ -308,7 +308,11 @@ func (a *Agent) process(input string, stageCtx *sdk.StageContext) (response stri // ContextPolicy: prune 工具调用后执行上下文裁剪(§13.8) if def := a.stageHost.ToolDef(tc.Name); def != nil && def.ContextPolicy == "prune" { if a.context != nil { - a.context.Prune(result, 20, a.docStore) + topK := a.maxContextSize - 1 + if topK < 1 { + topK = 1 + } + a.context.Prune(result, topK, a.docStore) } } diff --git a/internal/plugin/proc/corehandler.go b/internal/plugin/proc/corehandler.go index 49b6e28..a222020 100644 --- a/internal/plugin/proc/corehandler.go +++ b/internal/plugin/proc/corehandler.go @@ -154,28 +154,28 @@ func (h *coreHandler) Handle(method string, params json.RawMessage) (interface{} if err := unmarshal(params, &p); err != nil { return nil, err } - h.sdk.InjectText(p.Source, p.Channel, p.Text) + h.sdk.InjectText(p.Source, p.Channel, h.resolveText(p)) return nil, nil case MethodIOInjectInterrupt: var p injectParams if err := unmarshal(params, &p); err != nil { return nil, err } - h.sdk.InjectInterruptText(p.Source, p.Channel, p.Text) + h.sdk.InjectInterruptText(p.Source, p.Channel, h.resolveText(p)) return nil, nil case MethodIOInjectTextNoMem: var p injectParams if err := unmarshal(params, &p); err != nil { return nil, err } - h.sdk.InjectTextNoMemory(p.Source, p.Channel, p.Text) + h.sdk.InjectTextNoMemory(p.Source, p.Channel, h.resolveText(p)) return nil, nil case MethodIOInjectSync: var p injectParams if err := unmarshal(params, &p); err != nil { return nil, err } - return map[string]interface{}{"reply": h.sdk.InjectInputSync(p.Source, p.Channel, p.Text)}, nil + return map[string]interface{}{"reply": h.sdk.InjectInputSync(p.Source, p.Channel, h.resolveText(p))}, nil case MethodIOInjectMedia: var p injectMediaParams @@ -563,7 +563,7 @@ func (h *coreHandler) Handle(method string, params json.RawMessage) (interface{} // 否则使用内联 Text。兼容新旧两种协议。 func (h *coreHandler) resolveText(p injectParams) string { if !p.TextRef.IsZero() { - return string(p.TextRef.Slice(h.host.data)) + return string(h.host.unified.arenaRead(p.TextRef)) } return p.Text } @@ -610,6 +610,12 @@ func (h *coreHandler) cleanerProxy(scope, name string, enabled bool) (func(strin return nil, fmt.Errorf("%s %s 声明 Cleaner,但清洗回调通道未就绪", scope, name) } return func(text string) string { + // Cleaner 输入和输出共享同一 arena。串行覆盖完整往返,保证结果读完前 + // 不被其他调用重置;返回 string 会复制结果,随后即可回收整个临时区。 + h.host.arenaMu.Lock() + defer h.host.arenaMu.Unlock() + defer h.host.unified.arenaReset() + ref, err := h.host.unified.arenaWrite([]byte(text)) if err != nil { return text @@ -619,6 +625,9 @@ func (h *coreHandler) cleanerProxy(scope, name string, enabled bool) (func(strin return text } result := h.host.unified.arenaRead(resultRef) + if result == nil && !resultRef.IsZero() { + return text + } return string(result) }, nil } @@ -636,6 +645,11 @@ func (h *coreHandler) toolRegister(params json.RawMessage) (interface{}, error) if p.Name == "" { return nil, fmt.Errorf("tool.register: 缺少 name") } + switch p.Def.ContextPolicy { + case "", "none", "prune": + default: + return nil, fmt.Errorf("tool.register: context_policy 只允许 none/prune,实际 %q", p.Def.ContextPolicy) + } p.Def.Plugin = h.name // 函数本身不进 JSON;has_cleaner 只声明其存在,实际执行回到插件进程。 cleaner, err := h.cleanerProxy(CleanerScopeTool, p.Name, p.HasCleaner) diff --git a/internal/plugin/proc/host.go b/internal/plugin/proc/host.go index 4620aaf..78c7037 100644 --- a/internal/plugin/proc/host.go +++ b/internal/plugin/proc/host.go @@ -36,6 +36,7 @@ type Host struct { evtSubscriber EvtRingSubscriber locks *lockRegistry + arenaMu sync.Mutex // 动态 arena 跨进程调用的生命周期锁 stageMu sync.Mutex coordMu sync.Mutex coord *stageCoordinator diff --git a/internal/plugin/proc/toollane.go b/internal/plugin/proc/toollane.go index b365962..ca3428e 100644 --- a/internal/plugin/proc/toollane.go +++ b/internal/plugin/proc/toollane.go @@ -47,10 +47,11 @@ const ( toolFrameSize = 112 // 单帧总大小 // 帧状态 - toolFrameFree uint32 = 0 - toolFrameCalling uint32 = 1 // 内核写 input,插件待读 - toolFrameReading uint32 = 2 // 插件正在执行 - toolFrameReady uint32 = 3 // 插件写 output,内核待读 + toolFrameFree uint32 = 0 + toolFrameReserved uint32 = 1 // 内核已占帧,正在写入元数据 + toolFrameCalling uint32 = 2 // 内核写完 input,插件待读 + toolFrameReading uint32 = 3 // 插件正在执行 + toolFrameReady uint32 = 4 // 插件写完 output,内核待读 ) // ToolRing 头偏移(相对区域起始) @@ -116,10 +117,19 @@ func AttachToolRing(data []byte) (*ToolCallRing, error) { if magic != toolRingMagic { return nil, fmt.Errorf("tool ring: 魔数不匹配(0x%x)", magic) } + cap := getU32(data[trlOffCap:]) + frameSize := getU32(data[trlOffFrameSize:]) + version := getU32(data[trlOffVersion:]) + if version != toolRingVersion { + return nil, fmt.Errorf("tool ring: 版本不匹配(%d,期望 %d)", version, toolRingVersion) + } + if cap == 0 || frameSize != toolFrameSize || uint64(trlOffFrameBase)+uint64(cap)*uint64(frameSize) > uint64(len(data)) { + return nil, fmt.Errorf("tool ring: 布局非法(cap=%d frameSize=%d total=%d)", cap, frameSize, len(data)) + } return &ToolCallRing{ data: data, - cap: getU32(data[trlOffCap:]), - frameSize: getU32(data[trlOffFrameSize:]), + cap: cap, + frameSize: frameSize, framesBase: trlOffFrameBase, }, nil } @@ -128,13 +138,11 @@ func AttachToolRing(data []byte) (*ToolCallRing, error) { // // 环形扫描:从 writeIdx 开始,绕环一圈找 FREE 帧。全部忙时背压。 func (r *ToolCallRing) Reserve() (frameIdx uint32, err error) { - start := r.writeIdx.Load() + start := r.writeIdx.Add(1) - 1 for i := uint64(0); i < uint64(r.cap); i++ { idx := start + i fi := uint32(idx % uint64(r.cap)) - state := atomic.LoadUint32(r.statePtr(fi)) - if state == toolFrameFree { - r.writeIdx.Store(idx + 1) + if atomic.CompareAndSwapUint32(r.statePtr(fi), toolFrameFree, toolFrameReserved) { return fi, nil } } @@ -191,16 +199,20 @@ func (r *ToolCallRing) GetCallingFrame(frameIdx uint32) (name string, inputRef S return } -// SetReading 将帧状态设为 READING(插件开始执行)。 -func (r *ToolCallRing) SetReading(frameIdx uint32) { - atomic.StoreUint32(r.statePtr(frameIdx), toolFrameReading) +// SetReading 将 CALLING 帧原子转为 READING(插件开始执行)。 +func (r *ToolCallRing) SetReading(frameIdx uint32) bool { + return atomic.CompareAndSwapUint32(r.statePtr(frameIdx), toolFrameCalling, toolFrameReading) } -// SetReady 将帧状态设为 READY 并写入 output 描述符(插件执行完毕)。 -func (r *ToolCallRing) SetReady(frameIdx uint32, outputRef SharedRef) { +// SetReady 将 READING 帧设为 READY 并写入 output 描述符(插件执行完毕)。 +func (r *ToolCallRing) SetReady(frameIdx uint32, outputRef SharedRef) bool { + if atomic.LoadUint32(r.statePtr(frameIdx)) != toolFrameReading { + return false + } off := r.frameOff(frameIdx) packSharedRef(r.data[off+toolFrameOffOutput:], outputRef) atomic.StoreUint32(r.statePtr(frameIdx), toolFrameReady) + return true } // GetReadyFrame 读取帧的 requestID 和 output 描述符(内核消费 READY 帧)。 @@ -214,8 +226,8 @@ func (r *ToolCallRing) GetReadyFrame(frameIdx uint32) (reqID uint64, outputRef S // ReleaseFrame 回收帧为 FREE(内核消费完毕后调用)。 func (r *ToolCallRing) ReleaseFrame(frameIdx uint32) { off := r.frameOff(frameIdx) - // 清零帧内容 - for i := uint32(0); i < r.frameSize; i++ { + // state 必须最后发布 FREE;否则并发 Reserve 可能在其余字段尚未清零时复用帧。 + for i := uint32(4); i < r.frameSize; i++ { r.data[off+i] = 0 } atomic.StoreUint32(r.statePtr(frameIdx), toolFrameFree) diff --git a/internal/plugin/proc/toollane_test.go b/internal/plugin/proc/toollane_test.go index be8074a..c722513 100644 --- a/internal/plugin/proc/toollane_test.go +++ b/internal/plugin/proc/toollane_test.go @@ -1,6 +1,7 @@ package proc import ( + "sync" "testing" "unsafe" ) @@ -154,6 +155,56 @@ func TestSharedRef_PackUnpack(t *testing.T) { } } +func TestToolCallRing_ConcurrentReserveUnique(t *testing.T) { + data := make([]byte, int(trlOffFrameBase+toolRingCap*toolFrameSize)) + if err := InitToolRing(data); err != nil { + t.Fatal(err) + } + ring, err := AttachToolRing(data) + if err != nil { + t.Fatal(err) + } + + indices := make(chan uint32, toolRingCap) + var wg sync.WaitGroup + for i := uint32(0); i < toolRingCap; i++ { + wg.Add(1) + go func() { + defer wg.Done() + idx, reserveErr := ring.Reserve() + if reserveErr != nil { + t.Errorf("Reserve: %v", reserveErr) + return + } + indices <- idx + }() + } + wg.Wait() + close(indices) + + seen := make(map[uint32]bool, toolRingCap) + for idx := range indices { + if seen[idx] { + t.Fatalf("并发 Reserve 重复分配帧 %d", idx) + } + seen[idx] = true + } + if len(seen) != int(toolRingCap) { + t.Fatalf("唯一帧数=%d,期望 %d", len(seen), toolRingCap) + } +} + +func TestAttachToolRingRejectsInvalidLayout(t *testing.T) { + data := make([]byte, trlOffFrameBase) + putU32(data[trlOffMagic:], toolRingMagic) + putU32(data[trlOffVersion:], toolRingVersion) + putU32(data[trlOffCap:], toolRingCap) + putU32(data[trlOffFrameSize:], toolFrameSize) + if _, err := AttachToolRing(data); err == nil { + t.Fatal("越界布局应被拒绝") + } +} + func TestToolFrameLayoutAligned(t *testing.T) { // 确保帧布局字段偏移与内存布局一致(安全断言) var frame toolFrame diff --git a/internal/plugin/proc/unified.go b/internal/plugin/proc/unified.go index 3c1969c..0f3bf63 100644 --- a/internal/plugin/proc/unified.go +++ b/internal/plugin/proc/unified.go @@ -94,10 +94,11 @@ func (r SharedRef) IsZero() bool { // Slice 从共享内存中按 SharedRef 切片。data 必须是完整的 mmap 区域。 func (r SharedRef) Slice(data []byte) []byte { - if r.IsZero() || int(r.Offset)+int(r.Length) > len(data) { + end := uint64(r.Offset) + uint64(r.Length) + if r.IsZero() || end > uint64(len(data)) { return nil } - return data[r.Offset : r.Offset+r.Length] + return data[r.Offset:uint32(end)] } // unifiedRegion 统一共享内存区域的内核侧视图。 @@ -171,12 +172,22 @@ func attachUnifiedRegion(data []byte) (*unifiedRegion, error) { evtOff := getU32(data[sbOffEvtOff:]) evtSize := getU32(data[sbOffEvtSize:]) + capacity := getU32(data[sbOffCapacity:]) arenaOff := getU32(data[sbOffArenaOff:]) arenaCap := getU32(data[sbOffArenaCap:]) - if int(evtOff)+int(evtSize) > len(data) { + if capacity != uint32(len(data)) { + return nil, fmt.Errorf("unified: capacity 不匹配(header=%d mapped=%d)", capacity, len(data)) + } + if uint64(ctxOff)+uint64(ctxSize) > uint64(len(data)) { + return nil, fmt.Errorf("unified: StageContext 越界(off=%d size=%d total=%d)", ctxOff, ctxSize, len(data)) + } + if uint64(evtOff)+uint64(evtSize) > uint64(len(data)) { return nil, fmt.Errorf("unified: EvtRing 越界(off=%d size=%d total=%d)", evtOff, evtSize, len(data)) } + if uint64(arenaOff)+uint64(arenaCap) > uint64(len(data)) { + return nil, fmt.Errorf("unified: arena 越界(off=%d cap=%d total=%d)", arenaOff, arenaCap, len(data)) + } return &unifiedRegion{ data: data, @@ -248,16 +259,19 @@ func (r *unifiedRegion) arenaAlloc(n int) (uint32, error) { if n <= 0 { return 0, fmt.Errorf("arena: 非法分配长度 %d", n) } - used := atomic.LoadUint32(r.arenaUsed) - if used == 0 { - used = 1 + for { + used := atomic.LoadUint32(r.arenaUsed) + if used == 0 { + used = 1 + } + end := uint64(used) + uint64(n) + if end > uint64(r.arenaCap) { + return 0, fmt.Errorf("arena: 空间不足(需 %d,剩 %d)", n, uint64(r.arenaCap)-uint64(used)) + } + if atomic.CompareAndSwapUint32(r.arenaUsed, used, uint32(end)) { + return used, nil + } } - end := uint64(used) + uint64(n) - if end > uint64(r.arenaCap) { - return 0, fmt.Errorf("arena: 空间不足(需 %d,剩 %d)", n, uint64(r.arenaCap)-uint64(used)) - } - atomic.StoreUint32(r.arenaUsed, uint32(end)) - return used, nil } // arenaWrite 把 b 写入动态 arena 并返回 SharedRef。 @@ -275,8 +289,16 @@ func (r *unifiedRegion) arenaWrite(b []byte) (SharedRef, error) { return SharedRef{Offset: base, Length: uint32(len(b)), Generation: gen}, nil } -// arenaRead 按 SharedRef 读取数据。 +// arenaRead 按 SharedRef 读取数据。generation 不匹配或引用越出 arena 时拒绝。 func (r *unifiedRegion) arenaRead(ref SharedRef) []byte { + if ref.IsZero() || ref.Generation != uint32(r.generation()) { + return nil + } + end := uint64(ref.Offset) + uint64(ref.Length) + arenaEnd := uint64(r.arenaOff) + uint64(r.arenaCap) + if uint64(ref.Offset) < uint64(r.arenaOff)+1 || end > arenaEnd { + return nil + } return ref.Slice(r.data) }