Skip to content

Commit 3fc3b23

Browse files
Copilotphrocker
Copilot
and
phrocker
authored
Add pause/resume control for enterprise agent autonomous operations with auto-pause on chat (#32)
* Initial plan * Add pause/resume functionality to enterprise agent with state preservation Co-authored-by: phrocker <[email protected]> * Add tests and documentation for agent pause/resume functionality Co-authored-by: phrocker <[email protected]> * Address code review feedback: fix race condition, remove reflection, optimize ObjectMapper Co-authored-by: phrocker <[email protected]> * Final code review improvements: proper synchronization, error handling, and simplified pause check Co-authored-by: phrocker <[email protected]> * Fix test * Fix agent deployment * Auto-pause agent when chat session is established from dashboard Co-authored-by: phrocker <[email protected]> * fix --------- Co-authored-by: copilot-swe-agent[bot] <[email protected]> Co-authored-by: phrocker <[email protected]> Co-authored-by: Marc Parisi <[email protected]>
1 parent 2ed3005 commit 3fc3b23

File tree

8 files changed

+659
-13
lines changed

8 files changed

+659
-13
lines changed

‎docs/AGENT_PAUSE_RESUME.md‎

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
# Agent Chat Pause/Resume Functionality
2+
3+
## Overview
4+
5+
The enterprise agent now supports pausing and resuming autonomous operations via WebSocket chat commands. This allows users to:
6+
7+
1. Pause the agent's autonomous activities while preserving all state
8+
2. Modify the agent's context and planned operations while paused
9+
3. Resume operations from the exact point where they were paused
10+
11+
## WebSocket Commands
12+
13+
### Pause Agent
14+
Pauses all autonomous operations while preserving state including execution context and ZTAT tokens.
15+
16+
```json
17+
{
18+
"type": "pause-agent"
19+
}
20+
```
21+
22+
**Response:**
23+
```
24+
"Agent autonomous operations have been paused. All state has been preserved including execution context and ztats."
25+
```
26+
27+
### Resume Agent
28+
Resumes autonomous operations from the paused state.
29+
30+
```json
31+
{
32+
"type": "resume-agent"
33+
}
34+
```
35+
36+
**Response:**
37+
```
38+
"Agent autonomous operations have been resumed. Continuing from saved state."
39+
```
40+
41+
### Query Agent Status
42+
Check whether the agent is currently paused or running.
43+
44+
```json
45+
{
46+
"type": "agent-status"
47+
}
48+
```
49+
50+
**Response:**
51+
```
52+
"Agent status: PAUSED"
53+
or
54+
"Agent status: RUNNING"
55+
```
56+
57+
### Modify Context (Only when paused)
58+
Modify the agent's execution context or change planned operations while the agent is paused.
59+
60+
```json
61+
{
62+
"type": "modify-context",
63+
"contextKey": "customKey",
64+
"contextValue": "{\"data\": \"value\"}",
65+
"operation": "new_operation_name"
66+
}
67+
```
68+
69+
**Response (success):**
70+
```
71+
"Agent context has been modified. Changes will take effect when agent is resumed."
72+
```
73+
74+
**Response (error - agent not paused):**
75+
```
76+
"Cannot modify context while agent is running. Please pause the agent first."
77+
```
78+
79+
## Architecture
80+
81+
### Components Modified
82+
83+
1. **ChatAgent** (`enterprise-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/ChatAgent.java`)
84+
- Added pause/resume state management with thread-safe synchronization
85+
- Main agent loop checks pause state before executing operations
86+
- Provenance events track pause/resume actions
87+
88+
2. **ChatWSHandler** (`enterprise-agent/src/main/java/io/sentrius/agent/analysis/api/websocket/ChatWSHandler.java`)
89+
- Handles WebSocket commands: `pause-agent`, `resume-agent`, `agent-status`, `modify-context`
90+
- Tracks which session paused the agent
91+
92+
3. **WebSocky** (`enterprise-agent/src/main/java/io/sentrius/agent/analysis/model/WebSocky.java`)
93+
- Tracks pause state per session
94+
95+
4. **ProvenanceEvent** (`provenance-core/src/main/java/io/sentrius/sso/provenance/ProvenanceEvent.java`)
96+
- Added `AGENT_PAUSED` and `AGENT_RESUMED` event types for audit trail
97+
98+
## State Preservation
99+
100+
When the agent is paused, the following state is preserved:
101+
102+
- **Execution Context**: All messages, short-term memory, long-term memory
103+
- **ZTAT Tokens**: All zero-trust access tokens remain valid
104+
- **Agent Data**: All execution arguments and call parameters
105+
- **Verb Responses**: History of operations executed
106+
- **Communication Responses**: All LLM responses and conversations
107+
108+
## Use Cases
109+
110+
### 1. Emergency Stop
111+
Pause the agent immediately if it's performing unwanted actions:
112+
```json
113+
{"type": "pause-agent"}
114+
```
115+
116+
### 2. Context Modification
117+
Pause the agent, modify its behavior, then resume:
118+
```json
119+
{"type": "pause-agent"}
120+
{"type": "modify-context", "contextKey": "priority", "contextValue": "\"high\""}
121+
{"type": "resume-agent"}
122+
```
123+
124+
### 3. Workflow Adjustment
125+
Change the next operation the agent will perform:
126+
```json
127+
{"type": "pause-agent"}
128+
{"type": "modify-context", "operation": "list_ztat_requests"}
129+
{"type": "resume-agent"}
130+
```
131+
132+
### 4. Monitoring
133+
Check agent status during long-running operations:
134+
```json
135+
{"type": "agent-status"}
136+
```
137+
138+
## Thread Safety
139+
140+
The pause/resume functionality uses Java's synchronized blocks and `wait()`/`notifyAll()` mechanism to ensure thread-safe operation:
141+
142+
- The `paused` flag is volatile for visibility across threads
143+
- A dedicated `pauseLock` object synchronizes access to pause state
144+
- The main agent loop waits on the lock when paused
145+
- Resume operations notify waiting threads to continue
146+
147+
## Testing
148+
149+
Comprehensive unit tests validate the pause/resume functionality:
150+
151+
- Initial state (not paused)
152+
- Pause operation
153+
- Resume operation
154+
- Multiple pause calls (idempotent)
155+
- Resume without pause (no-op)
156+
- State retrieval
157+
- Shutdown while paused
158+
159+
Run tests:
160+
```bash
161+
mvn test -pl enterprise-agent -Dtest=ChatAgentPauseTest
162+
```
163+
164+
## Security Considerations
165+
166+
1. **Authorization**: Only authorized users should be able to pause/resume agents
167+
2. **Audit Trail**: All pause/resume operations are logged via provenance events
168+
3. **State Protection**: Context can only be modified when the agent is paused
169+
4. **ZTAT Preservation**: Zero-trust tokens remain valid during pause, maintaining security posture
170+
171+
## Future Enhancements
172+
173+
Potential improvements for future iterations:
174+
175+
- Multi-agent coordination: Pause all related agents
176+
- Scheduled pause/resume: Time-based automation
177+
- Conditional pause: Automatic pause on specific conditions
178+
- State snapshots: Save/restore multiple pause points
179+
- Rollback capability: Revert to previous states

‎enterprise-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/ChatAgent.java‎

Lines changed: 124 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ public class ChatAgent extends BaseEnterpriseAgent {
5050
final ChatVerbs chatVerbs;
5151

5252
private volatile boolean running = true;
53+
private volatile boolean paused = false;
54+
private final Object pauseLock = new Object();
5355
private Thread workerThread;
5456

5557
private AgentExecution agentExecution;
@@ -181,6 +183,23 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
181183
List<VerbResponse> verbResponses = new ArrayList<>();
182184
while(running) {
183185

186+
// Check if agent is paused if autonomous mode
187+
if (null != agentConfigOptions.getType() && agentConfigOptions.getType().equalsIgnoreCase("chat" +
188+
"-autonomous")) {
189+
synchronized (pauseLock) {
190+
while (paused) {
191+
try {
192+
log.info("Agent paused, waiting for resume command...");
193+
pauseLock.wait();
194+
log.info("Agent resumed from pause.");
195+
} catch (InterruptedException e) {
196+
Thread.currentThread().interrupt();
197+
log.warn("Agent interrupted while paused");
198+
break;
199+
}
200+
}
201+
}
202+
}
184203

185204
try {
186205

@@ -277,16 +296,117 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
277296

278297
}
279298

299+
/**
300+
* Pause the agent's autonomous operations.
301+
* Preserves the current state including execution context and ztats.
302+
*/
303+
public void pauseAgent() {
304+
synchronized (pauseLock) {
305+
if (!paused) {
306+
paused = true;
307+
log.info("Agent paused - state preserved");
308+
309+
// Submit provenance event for pause
310+
try {
311+
agentClientService.submitProvenance(
312+
agentExecution,
313+
io.sentrius.sso.provenance.ProvenanceEvent.builder()
314+
.eventType(io.sentrius.sso.provenance.ProvenanceEvent.EventType.AGENT_PAUSED)
315+
.actor(agentExecution.getUser().getUsername())
316+
.triggeringUser(agentExecution.getUser().getUsername())
317+
.outputSummary("Agent autonomous operations paused by user")
318+
.build()
319+
);
320+
} catch (Exception e) {
321+
log.error("Failed to submit pause provenance event", e);
322+
}
323+
}
324+
}
325+
}
326+
327+
/**
328+
* Resume the agent's autonomous operations.
329+
* Continues from the previously saved state.
330+
*/
331+
public void resumeAgent() {
332+
synchronized (pauseLock) {
333+
if (paused) {
334+
paused = false;
335+
pauseLock.notifyAll();
336+
log.info("Agent resumed - continuing operations");
337+
338+
// Submit provenance event for resume
339+
try {
340+
agentClientService.submitProvenance(
341+
agentExecution,
342+
io.sentrius.sso.provenance.ProvenanceEvent.builder()
343+
.eventType(io.sentrius.sso.provenance.ProvenanceEvent.EventType.AGENT_RESUMED)
344+
.actor(agentExecution.getUser().getUsername())
345+
.triggeringUser(agentExecution.getUser().getUsername())
346+
.outputSummary("Agent autonomous operations resumed by user")
347+
.build()
348+
);
349+
} catch (Exception e) {
350+
log.error("Failed to submit resume provenance event", e);
351+
}
352+
}
353+
}
354+
}
355+
356+
/**
357+
* Check if the agent is currently paused.
358+
*/
359+
public boolean isPaused() {
360+
return paused;
361+
}
362+
363+
/**
364+
* Execute a context modification if the agent is paused.
365+
* This method ensures thread-safe modification of agent context.
366+
*
367+
* @param modifier The runnable that performs the context modification
368+
* @return true if the modification was performed, false if agent is not paused
369+
*/
370+
public boolean modifyContextIfPaused(Runnable modifier) {
371+
synchronized (pauseLock) {
372+
if (paused) {
373+
modifier.run();
374+
return true;
375+
}
376+
return false;
377+
}
378+
}
379+
380+
/**
381+
* Get the current agent execution context.
382+
* This includes all state, messages, and execution data.
383+
*/
384+
public AgentExecution getAgentExecution() {
385+
return agentExecution;
386+
}
387+
388+
/**
389+
* Set the agent execution context.
390+
* Public method primarily for testing purposes.
391+
* Should not be used in production code.
392+
*/
393+
public void setAgentExecution(AgentExecution agentExecution) {
394+
this.agentExecution = agentExecution;
395+
}
396+
280397
@PreDestroy
281398
public void shutdown() {
282399
log.info("Shutting down ChatAgent...");
283400
running = false;
401+
402+
// Wake up any paused threads
403+
synchronized (pauseLock) {
404+
paused = false;
405+
pauseLock.notifyAll();
406+
}
407+
284408
if (workerThread != null) {
285409
workerThread.interrupt();
286410
}
287411
}
288-
289-
public AgentExecution getAgentExecution() {
290-
return agentExecution;
291-
}
292412
}

‎enterprise-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AgentVerbs.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -736,8 +736,7 @@ public JsonNode getAgentExecutionStatus(AgentExecution execution, AgentExecution
736736

737737
@Verb(name = "create_agent", returnType = AgentExecutionContextDTO.class, description = "Creates an agent who has the " +
738738
"context. a previously defined contextId is required. previously defined endpoints can be used to build a " +
739-
"trust policy. must call create_agent_context before this verb. agent type is chat, chat-autonomous, or " +
740-
"autonomous. chat is chat only, chat-autonomous is chat and autonomous. determine based on workload.",
739+
"trust policy. must call create_agent_context before this verb. agent type is chat or chat-autonomous. chat is chat only, chat-autonomous is chat and autonomous. determine based on workload.",
741740
exampleJson = "{ \"agentName\": \"agentName\", \"agentType\": \"agentType\" }",
742741
requiresTokenManagement = true )
743742
public ObjectNode createAgent(AgentExecution execution, AgentExecutionContextDTO context)
@@ -848,12 +847,18 @@ public ObjectNode getEndpointsLike(AgentExecution execution,
848847
ArrayNode endpoints = JsonUtil.MAPPER.createArrayNode();
849848
for(JsonNode node : parsedQuery) {
850849
if (!node.isTextual()) {
851-
throw new IllegalArgumentException("All items in 'endpoints_like' must be strings");
850+
if (node.has("arg1")){
851+
node = node.get("arg1");
852+
if (!node.isTextual()) {
853+
throw new IllegalArgumentException("All items in 'endpoints_like' must be strings");
854+
}
855+
}
852856
}
853857
var endpointList = endpointSearcher.getEndpointsLike(execution, node.asText());
858+
JsonNode finalNode = node;
854859
endpointList.forEach(endpoint -> {
855860
ObjectNode endpointNode = JsonUtil.MAPPER.createObjectNode();
856-
endpointNode.put("name", node.asText());
861+
endpointNode.put("name", finalNode.asText());
857862
endpointNode.put("method", endpoint.getHttpMethod());
858863
endpointNode.put("endpoint", endpoint.getPath());
859864
endpoints.add(endpointNode);

0 commit comments

Comments
 (0)