Skip to content

Commit b554b1a

Browse files
EhabYEhabY
committed
Add agent state machine
1 parent b1785ed commit b554b1a

File tree

2 files changed

+146
-79
lines changed

2 files changed

+146
-79
lines changed

‎src/api/workspace.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,6 @@ export async function waitForBuild(
138138

139139
/**
140140
* Streams agent logs to the emitter in real-time.
141-
* Fetches existing logs and subscribes to new logs via websocket.
142141
* Returns the websocket and a completion promise that rejects on error.
143142
*/
144143
export async function streamAgentLogs(

‎src/remote/remote.ts‎

Lines changed: 146 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -539,81 +539,13 @@ export class Remote {
539539
const inbox = await Inbox.create(workspace, workspaceClient, this.logger);
540540
disposables.push(inbox);
541541

542-
// Wait for the agent to connect.
543-
if (agent.status === "connecting") {
544-
this.logger.info(`Waiting for ${workspaceName}/${agent.name}...`);
545-
const updatedAgent = await this.waitForAgentWithProgress(
546-
monitor,
547-
agent,
548-
`Waiting for agent ${agent.name} to connect...`,
549-
(foundAgent) => foundAgent.status !== "connecting",
550-
);
551-
agent = updatedAgent;
552-
this.logger.info(`Agent ${agent.name} status is now`, agent.status);
553-
}
554-
555-
// Make sure the agent is connected.
556-
if (agent.status !== "connected") {
557-
const result = await this.vscodeProposed.window.showErrorMessage(
558-
`${workspaceName}/${agent.name} ${agent.status}`,
559-
{
560-
useCustom: true,
561-
modal: true,
562-
detail: `The ${agent.name} agent failed to connect. Try restarting your workspace.`,
563-
},
564-
);
565-
if (!result) {
566-
await this.closeRemote();
567-
return;
568-
}
569-
await this.reloadWindow();
570-
return;
571-
}
572-
573-
if (agent.lifecycle_state === "starting") {
574-
const isBlocking = agent.scripts.some(
575-
(script) => script.start_blocks_login,
576-
);
577-
if (isBlocking) {
578-
this.logger.info(
579-
`Waiting for ${workspaceName}/${agent.name} startup...`,
580-
);
581-
582-
let terminalSession: TerminalSession | undefined;
583-
let socket: OneWayWebSocket<WorkspaceAgentLog[]> | undefined;
584-
try {
585-
terminalSession = new TerminalSession("Agent Log");
586-
const { socket: agentSocket, completion: logsCompletion } =
587-
await streamAgentLogs(
588-
workspaceClient,
589-
terminalSession.writeEmitter,
590-
agent,
591-
);
592-
socket = agentSocket;
593-
594-
const agentStatePromise = this.waitForAgentWithProgress(
595-
monitor,
596-
agent,
597-
`Waiting for agent ${agent.name} startup scripts...`,
598-
(foundAgent) => foundAgent.lifecycle_state !== "starting",
599-
);
600-
601-
// Race between logs completion and agent state change
602-
const updatedAgent = await Promise.race([
603-
agentStatePromise,
604-
logsCompletion.then(() => agentStatePromise),
605-
]);
606-
agent = updatedAgent;
607-
this.logger.info(
608-
`Agent ${agent.name} lifecycle state is now`,
609-
agent.lifecycle_state,
610-
);
611-
} finally {
612-
terminalSession?.dispose();
613-
socket?.close();
614-
}
615-
}
616-
}
542+
// Ensure agent is ready by handling all status and lifecycle states
543+
agent = await this.ensureAgentReady(
544+
monitor,
545+
agent,
546+
workspaceName,
547+
workspaceClient,
548+
);
617549

618550
const logDir = this.getLogDir(featureSet);
619551

@@ -740,6 +672,139 @@ export class Remote {
740672
return ` --log-dir ${escapeCommandArg(logDir)} -v`;
741673
}
742674

675+
/**
676+
* Ensures agent is ready to connect by handling all status and lifecycle states.
677+
* Throws an error if the agent cannot be made ready.
678+
*/
679+
private async ensureAgentReady(
680+
monitor: WorkspaceMonitor,
681+
agent: WorkspaceAgent,
682+
workspaceName: string,
683+
workspaceClient: CoderApi,
684+
): Promise<WorkspaceAgent> {
685+
let currentAgent = agent;
686+
687+
while (
688+
currentAgent.status !== "connected" ||
689+
currentAgent.lifecycle_state !== "ready"
690+
) {
691+
switch (currentAgent.status) {
692+
case "connecting":
693+
this.logger.info(`Waiting for agent ${currentAgent.name}...`);
694+
currentAgent = await this.waitForAgentWithProgress(
695+
monitor,
696+
currentAgent,
697+
`Waiting for agent ${currentAgent.name} to connect...`,
698+
(foundAgent) => foundAgent.status !== "connecting",
699+
);
700+
this.logger.info(
701+
`Agent ${currentAgent.name} status is now`,
702+
currentAgent.status,
703+
);
704+
continue;
705+
706+
case "connected":
707+
// Agent connected, now handle lifecycle state
708+
break;
709+
710+
case "disconnected":
711+
case "timeout":
712+
throw new Error(
713+
`${workspaceName}/${currentAgent.name} ${currentAgent.status}`,
714+
);
715+
716+
default:
717+
throw new Error(
718+
`${workspaceName}/${currentAgent.name} unknown status: ${currentAgent.status}`,
719+
);
720+
}
721+
722+
// Handle agent lifecycle state (only when status is "connected")
723+
switch (currentAgent.lifecycle_state) {
724+
case "ready":
725+
return currentAgent;
726+
727+
case "starting": {
728+
const isBlocking = currentAgent.scripts.some(
729+
(script) => script.start_blocks_login,
730+
);
731+
if (!isBlocking) {
732+
return currentAgent;
733+
}
734+
735+
const logMsg = `Waiting for agent ${currentAgent.name} startup scripts...`;
736+
this.logger.info(logMsg);
737+
738+
let terminalSession: TerminalSession | undefined;
739+
let socket: OneWayWebSocket<WorkspaceAgentLog[]> | undefined;
740+
try {
741+
terminalSession = new TerminalSession("Agent Log");
742+
const { socket: agentSocket, completion: logsCompletion } =
743+
await streamAgentLogs(
744+
workspaceClient,
745+
terminalSession.writeEmitter,
746+
currentAgent,
747+
);
748+
socket = agentSocket;
749+
750+
const agentStatePromise = this.waitForAgentWithProgress(
751+
monitor,
752+
currentAgent,
753+
logMsg,
754+
(foundAgent) => foundAgent.lifecycle_state !== "starting",
755+
);
756+
757+
// Race between logs completion and agent state change
758+
currentAgent = await Promise.race([
759+
agentStatePromise,
760+
logsCompletion.then(() => agentStatePromise),
761+
]);
762+
this.logger.info(
763+
`Agent ${currentAgent.name} lifecycle state is now`,
764+
currentAgent.lifecycle_state,
765+
);
766+
} finally {
767+
terminalSession?.dispose();
768+
socket?.close();
769+
}
770+
continue;
771+
}
772+
773+
case "created":
774+
this.logger.info(
775+
`Waiting for ${workspaceName}/${currentAgent.name} to start...`,
776+
);
777+
currentAgent = await this.waitForAgentWithProgress(
778+
monitor,
779+
currentAgent,
780+
`Waiting for agent ${currentAgent.name} to start...`,
781+
(foundAgent) => foundAgent.lifecycle_state !== "created",
782+
);
783+
this.logger.info(
784+
`Agent ${currentAgent.name} lifecycle state is now`,
785+
currentAgent.lifecycle_state,
786+
);
787+
continue;
788+
789+
case "off":
790+
case "start_error":
791+
case "start_timeout":
792+
case "shutdown_error":
793+
case "shutdown_timeout":
794+
case "shutting_down":
795+
throw new Error(
796+
`${workspaceName}/${currentAgent.name} lifecycle state: ${currentAgent.lifecycle_state}`,
797+
);
798+
799+
default:
800+
throw new Error(
801+
`${workspaceName}/${currentAgent.name} unknown lifecycle state: ${currentAgent.lifecycle_state}`,
802+
);
803+
}
804+
}
805+
return currentAgent;
806+
}
807+
743808
/**
744809
* Waits for an agent condition with progress notification.
745810
*/
@@ -749,7 +814,7 @@ export class Remote {
749814
progressTitle: string,
750815
checker: (agent: WorkspaceAgent) => boolean,
751816
): Promise<WorkspaceAgent> {
752-
const foundAgent = await vscode.window.withProgress(
817+
const foundAgent = await this.vscodeProposed.window.withProgress(
753818
{
754819
title: progressTitle,
755820
location: vscode.ProgressLocation.Notification,
@@ -770,6 +835,10 @@ export class Remote {
770835
): Promise<WorkspaceAgent> {
771836
return new Promise<WorkspaceAgent>((resolve, reject) => {
772837
const updateEvent = monitor.onChange.event((workspace) => {
838+
if (workspace.latest_build.status !== "running") {
839+
const workspaceName = createWorkspaceIdentifier(workspace);
840+
reject(new Error(`Workspace ${workspaceName} is not running.`));
841+
}
773842
try {
774843
const agents = extractAgents(workspace.latest_build.resources);
775844
const foundAgent = agents.find((a) => a.id === agent.id);
@@ -778,8 +847,7 @@ export class Remote {
778847
`Agent ${agent.name} not found in workspace resources`,
779848
);
780849
}
781-
const result = checker(foundAgent);
782-
if (result !== undefined) {
850+
if (checker(foundAgent)) {
783851
updateEvent.dispose();
784852
resolve(foundAgent);
785853
}

0 commit comments

Comments
 (0)