返回 DeepSeek-Reasonix
workspace_lease_test.go
根目录 / internal / agent / workspace_lease_test.go
1 package agent
2
3 import (
4 "context"
5 "encoding/json"
6 "sync/atomic"
7 "testing"
8 "time"
9
10 "reasonix/internal/event"
11 "reasonix/internal/evidence"
12 "reasonix/internal/provider"
13 "reasonix/internal/tool"
14 "reasonix/internal/workspacelease"
15 )
16
17 type workspaceLeaseTestTool struct {
18 name string
19 readOnly bool
20 calls atomic.Int32
21 }
22
23 type workspaceLeaseTestHooks struct{ preCalls atomic.Int32 }
24 type workspaceLeaseDenyGate struct{}
25
26 func (workspaceLeaseDenyGate) Check(context.Context, string, json.RawMessage, bool) (bool, string, error) {
27 return false, "test denial", nil
28 }
29
30 func (h *workspaceLeaseTestHooks) PreToolUse(context.Context, string, json.RawMessage) (bool, string) {
31 h.preCalls.Add(1)
32 return false, ""
33 }
34 func (*workspaceLeaseTestHooks) PostToolUse(context.Context, string, json.RawMessage, string) {}
35 func (*workspaceLeaseTestHooks) PostToolUseFailure(context.Context, string, json.RawMessage, string, error) {
36 }
37 func (*workspaceLeaseTestHooks) PostLLMCall(_ context.Context, reasoning string, _ int) string {
38 return reasoning
39 }
40 func (*workspaceLeaseTestHooks) HasPostLLMCall() bool { return false }
41 func (*workspaceLeaseTestHooks) SubagentStop(context.Context, string) {}
42 func (*workspaceLeaseTestHooks) PreCompact(context.Context, string) string { return "" }
43
44 func (t *workspaceLeaseTestTool) Name() string { return t.name }
45 func (t *workspaceLeaseTestTool) Description() string { return t.name }
46 func (t *workspaceLeaseTestTool) Schema() json.RawMessage {
47 return json.RawMessage(`{"type":"object"}`)
48 }
49 func (t *workspaceLeaseTestTool) ReadOnly() bool { return t.readOnly }
50 func (t *workspaceLeaseTestTool) Execute(context.Context, json.RawMessage) (string, error) {
51 t.calls.Add(1)
52 return "ok", nil
53 }
54
55 func deliveryLeaseTestAgent(t *testing.T, owner *workspacelease.Owner, tools ...tool.Tool) *Agent {
56 t.Helper()
57 reg := tool.NewRegistry()
58 for _, candidate := range tools {
59 reg.Add(candidate)
60 }
61 a := New(nil, reg, NewSession(""), Options{DeliveryProfile: true, WorkspaceLease: owner}, event.Discard)
62 a.deliveryCriteriaEstablished = true
63 a.setTodoState([]evidence.TodoItem{{Content: "mutate", Status: "in_progress"}})
64 return a
65 }
66
67 func TestDeliveryWriterWaitsBeforeToolExecutionButReaderDoesNot(t *testing.T) {
68 root, locks := t.TempDir(), t.TempDir()
69 first, err := workspacelease.New(root, locks, nil)
70 if err != nil {
71 t.Fatal(err)
72 }
73 second, err := workspacelease.New(root, locks, nil)
74 if err != nil {
75 t.Fatal(err)
76 }
77 first.BeginRun()
78 if err := first.AcquireWrite(context.Background()); err != nil {
79 t.Fatal(err)
80 }
81 defer first.EndRun()
82
83 reader := &workspaceLeaseTestTool{name: "lease_reader", readOnly: true}
84 writer := &workspaceLeaseTestTool{name: "lease_writer", readOnly: false}
85 a := deliveryLeaseTestAgent(t, second, reader, writer)
86 second.BeginRun()
87 defer second.EndRun()
88
89 if outcome := a.executeOne(context.Background(), providerToolCall("read", reader.Name())); outcome.errMsg != "" {
90 t.Fatalf("reader was blocked by another Delivery writer: %+v", outcome)
91 }
92 if got := reader.calls.Load(); got != 1 {
93 t.Fatalf("reader calls = %d, want 1", got)
94 }
95
96 hooks := &workspaceLeaseTestHooks{}
97 a.hooks = hooks
98 ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
99 defer cancel()
100 outcome := a.executeOne(ctx, providerToolCall("write", writer.Name()))
101 if !outcome.blocked || outcome.errMsg != "blocked: workspace write lease unavailable" {
102 t.Fatalf("writer outcome = %+v, want lease block", outcome)
103 }
104 if got := writer.calls.Load(); got != 0 {
105 t.Fatalf("writer executed %d times before lease acquisition", got)
106 }
107 if got := hooks.preCalls.Load(); got != 0 {
108 t.Fatalf("PreToolUse ran %d times before lease acquisition", got)
109 }
110 }
111
112 func TestDeniedDeliveryWriterDoesNotAcquireWorkspaceLease(t *testing.T) {
113 root, locks := t.TempDir(), t.TempDir()
114 deniedOwner, _ := workspacelease.New(root, locks, nil)
115 probeOwner, _ := workspacelease.New(root, locks, nil)
116 writer := &workspaceLeaseTestTool{name: "denied_writer"}
117 a := deliveryLeaseTestAgent(t, deniedOwner, writer)
118 a.gate = workspaceLeaseDenyGate{}
119 deniedOwner.BeginRun()
120 outcome := a.executeOne(context.Background(), providerToolCall("write", writer.Name()))
121 deniedOwner.EndRun()
122 if !outcome.blocked || outcome.errMsg != "blocked by permission policy" {
123 t.Fatalf("denied outcome = %+v", outcome)
124 }
125 if writer.calls.Load() != 0 {
126 t.Fatal("denied writer executed")
127 }
128 probeOwner.BeginRun()
129 ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
130 defer cancel()
131 if err := probeOwner.AcquireWrite(ctx); err != nil {
132 t.Fatalf("permission denial leaked workspace lease: %v", err)
133 }
134 probeOwner.EndRun()
135 }
136
137 func providerToolCall(id, name string) provider.ToolCall {
138 return provider.ToolCall{ID: id, Name: name, Arguments: `{}`}
139 }
140
140 lines GO