@agentmeshhq/agent 0.1.4 → 0.1.6

package/src/core/daemon.ts

@@ -1,15 +1,22 @@
-import { createSession, destroySession, sessionExists, getSessionName } from "./tmux.js";
-import { AgentWebSocket } from "./websocket.js";
-import { Heartbeat } from "./heartbeat.js";
-import { registerAgent, checkInbox } from "./registry.js";
-import { injectStartupMessage, handleWebSocketEvent } from "./injector.js";
 import {
   addAgentToState,
-  removeAgentFromState,
-  loadConfig,
   getAgentState,
+  loadConfig,
+  removeAgentFromState,
+  updateAgentInState,
 } from "../config/loader.js";
-import type { Config, AgentConfig } from "../config/schema.js";
+import type { AgentConfig, Config } from "../config/schema.js";
+import { Heartbeat } from "./heartbeat.js";
+import { handleWebSocketEvent, injectStartupMessage } from "./injector.js";
+import { checkInbox, registerAgent } from "./registry.js";
+import {
+  createSession,
+  destroySession,
+  getSessionName,
+  sessionExists,
+  updateSessionEnvironment,
+} from "./tmux.js";
+import { AgentWebSocket } from "./websocket.js";
 
 export interface DaemonOptions {
   name: string;
@@ -32,9 +39,7 @@ export class AgentDaemon {
   constructor(options: DaemonOptions) {
     const config = loadConfig();
     if (!config) {
-      throw new Error(
-        "No config found. Run 'agentmesh init' first."
-      );
+      throw new Error("No config found. Run 'agentmesh init' first.");
     }
 
     this.config = config;
@@ -78,7 +83,7 @@ export class AgentDaemon {
       const created = createSession(
         this.agentName,
         this.agentConfig.command,
-        this.agentConfig.workdir
+        this.agentConfig.workdir,
       );
 
       if (!created) {
@@ -106,6 +111,13 @@ export class AgentDaemon {
 
     console.log(`Registered as: ${this.agentId}`);
 
+    // Inject environment variables into tmux session
+    console.log("Injecting environment variables...");
+    updateSessionEnvironment(this.agentName, {
+      AGENT_TOKEN: this.token,
+      AGENTMESH_AGENT_ID: this.agentId,
+    });
+
     // Save state
     addAgentToState({
       name: this.agentName,
@@ -116,15 +128,53 @@ export class AgentDaemon {
       token: this.token,
     });
 
-    // Start heartbeat
+    // Start heartbeat with auto-refresh
     console.log("Starting heartbeat...");
     this.heartbeat = new Heartbeat({
       url: this.config.hubUrl,
       token: this.token,
       intervalMs: 30000,
+      agentName: this.agentName,
+      agentId: this.agentId,
+      apiKey: this.config.apiKey,
+      workspace: this.config.workspace,
       onError: (error) => {
         console.error("Heartbeat error:", error.message);
       },
+      onTokenRefresh: (newToken) => {
+        this.token = newToken;
+        // Update state file
+        updateAgentInState(this.agentName, { token: newToken });
+        // Update tmux environment
+        updateSessionEnvironment(this.agentName, {
+          AGENT_TOKEN: newToken,
+          AGENTMESH_AGENT_ID: this.agentId!,
+        });
+        // Reconnect WebSocket with new token
+        if (this.ws) {
+          this.ws.disconnect();
+          const wsUrl = this.config.hubUrl
+            .replace("https://", "wss://")
+            .replace("http://", "ws://");
+          this.ws = new AgentWebSocket({
+            url: `${wsUrl}/ws/v1`,
+            token: newToken,
+            onMessage: (event) => {
+              handleWebSocketEvent(this.agentName, event);
+            },
+            onConnect: () => {
+              console.log("WebSocket reconnected with new token");
+            },
+            onDisconnect: () => {
+              console.log("WebSocket disconnected");
+            },
+            onError: (error) => {
+              console.error("WebSocket error:", error.message);
+            },
+          });
+          this.ws.connect();
+        }
+      },
     });
     this.heartbeat.start();
 
@@ -153,11 +203,7 @@ export class AgentDaemon {
     // Check inbox and auto-nudge
     console.log("Checking inbox...");
     try {
-      const inboxItems = await checkInbox(
-        this.config.hubUrl,
-        this.config.workspace,
-        this.token
-      );
+      const inboxItems = await checkInbox(this.config.hubUrl, this.config.workspace, this.token);
       injectStartupMessage(this.agentName, inboxItems.length);
     } catch (error) {
       console.error("Failed to check inbox:", error);

package/src/core/heartbeat.ts

@@ -1,16 +1,28 @@
+import { getTokenTimeRemaining, isTokenExpiringSoon } from "../utils/jwt.js";
+
+// Refresh token when less than 24 hours remain
+const TOKEN_REFRESH_THRESHOLD_MS = 24 * 60 * 60 * 1000;
+
 export interface HeartbeatConfig {
   url: string;
   token: string;
   intervalMs: number;
+  agentName: string;
+  agentId: string;
+  apiKey: string;
+  workspace: string;
   onError?: (error: Error) => void;
+  onTokenRefresh?: (newToken: string) => void;
 }
 
 export class Heartbeat {
   private config: HeartbeatConfig;
+  private currentToken: string;
   private intervalId: NodeJS.Timeout | null = null;
 
   constructor(config: HeartbeatConfig) {
     this.config = config;
+    this.currentToken = config.token;
   }
 
   start(): void {
@@ -34,22 +46,71 @@ export class Heartbeat {
     }
   }
 
+  getToken(): string {
+    return this.currentToken;
+  }
+
   private async sendHeartbeat(): Promise<void> {
     try {
+      // Check if token needs refresh before sending heartbeat
+      if (isTokenExpiringSoon(this.currentToken, TOKEN_REFRESH_THRESHOLD_MS)) {
+        const remaining = getTokenTimeRemaining(this.currentToken);
+        const hours = Math.floor(remaining / (1000 * 60 * 60));
+        console.log(`Token expiring in ${hours} hours, refreshing...`);
+        await this.refreshToken();
+      }
+
       const response = await fetch(`${this.config.url}/api/v1/agents/heartbeat`, {
         method: "POST",
         headers: {
-          Authorization: `Bearer ${this.config.token}`,
+          Authorization: `Bearer ${this.currentToken}`,
           "Content-Type": "application/json",
         },
         body: JSON.stringify({}),
       });
 
       if (!response.ok) {
+        // If unauthorized, try to refresh token
+        if (response.status === 401) {
+          console.log("Token rejected, attempting refresh...");
+          await this.refreshToken();
+          return;
+        }
         throw new Error(`Heartbeat failed: ${response.status}`);
       }
     } catch (error) {
       this.config.onError?.(error as Error);
     }
   }
+
+  private async refreshToken(): Promise<void> {
+    try {
+      const response = await fetch(`${this.config.url}/api/v1/agents/register`, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          "x-agentmesh-secret": this.config.apiKey,
+        },
+        body: JSON.stringify({
+          agent_id: this.config.agentId,
+          workspace: this.config.workspace,
+          display_name: this.config.agentName,
+          model: "claude-sonnet-4",
+        }),
+      });
+
+      if (!response.ok) {
+        throw new Error(`Token refresh failed: ${response.status}`);
+      }
+
+      const data = (await response.json()) as { token: string };
+      this.currentToken = data.token;
+
+      console.log("Token refreshed successfully");
+      this.config.onTokenRefresh?.(data.token);
+    } catch (error) {
+      console.error("Failed to refresh token:", (error as Error).message);
+      this.config.onError?.(error as Error);
+    }
+  }
 }

package/src/core/tmux.ts

@@ -1,4 +1,4 @@
-import { execSync, spawn, type ChildProcess } from "node:child_process";
+import { type ChildProcess, execSync, spawn } from "node:child_process";
 
 const SESSION_PREFIX = "agentmesh-";
 
@@ -15,10 +15,16 @@ export function sessionExists(sessionName: string): boolean {
   }
 }
 
+export interface SessionEnv {
+  AGENT_TOKEN?: string;
+  AGENTMESH_AGENT_ID?: string;
+}
+
 export function createSession(
   agentName: string,
   command: string,
-  workdir?: string
+  workdir?: string,
+  env?: SessionEnv,
 ): boolean {
   const sessionName = getSessionName(agentName);
 
@@ -37,6 +43,12 @@ export function createSession(
     args.push(command);
 
     execSync(`tmux ${args.join(" ")}`);
+
+    // Inject environment variables after session creation
+    if (env) {
+      setSessionEnvironment(sessionName, env);
+    }
+
     return true;
   } catch (error) {
     console.error(`Failed to create tmux session: ${error}`);
@@ -44,6 +56,34 @@ export function createSession(
   }
 }
 
+export function setSessionEnvironment(sessionName: string, env: SessionEnv): boolean {
+  try {
+    if (env.AGENT_TOKEN) {
+      execSync(`tmux set-environment -t "${sessionName}" AGENT_TOKEN "${env.AGENT_TOKEN}"`);
+    }
+    if (env.AGENTMESH_AGENT_ID) {
+      execSync(
+        `tmux set-environment -t "${sessionName}" AGENTMESH_AGENT_ID "${env.AGENTMESH_AGENT_ID}"`,
+      );
+    }
+    return true;
+  } catch (error) {
+    console.error(`Failed to set session environment: ${error}`);
+    return false;
+  }
+}
+
+export function updateSessionEnvironment(agentName: string, env: SessionEnv): boolean {
+  const sessionName = getSessionName(agentName);
+
+  if (!sessionExists(sessionName)) {
+    console.error(`Session ${sessionName} does not exist`);
+    return false;
+  }
+
+  return setSessionEnvironment(sessionName, env);
+}
+
 export function destroySession(agentName: string): boolean {
   const sessionName = getSessionName(agentName);
 
@@ -119,9 +159,7 @@ export function listSessions(): string[] {
   }
 }
 
-export function getSessionInfo(
-  agentName: string
-): { exists: boolean; command?: string } {
+export function getSessionInfo(agentName: string): { exists: boolean; command?: string } {
   const sessionName = getSessionName(agentName);
 
   if (!sessionExists(sessionName)) {
@@ -129,10 +167,9 @@ export function getSessionInfo(
   }
 
   try {
-    const command = execSync(
-      `tmux list-panes -t "${sessionName}" -F "#{pane_current_command}"`,
-      { encoding: "utf-8" }
-    ).trim();
+    const command = execSync(`tmux list-panes -t "${sessionName}" -F "#{pane_current_command}"`, {
+      encoding: "utf-8",
+    }).trim();
 
     return { exists: true, command };
   } catch {

package/src/utils/jwt.ts

@@ -0,0 +1,87 @@
+/**
+ * JWT utilities for token management
+ * Note: We only decode the payload, we don't verify signatures here
+ * (verification happens on the hub side)
+ */
+
+export interface TokenPayload {
+  sub: string; // Agent ID
+  actorType: string; // "agent"
+  workspaceScopes: string[];
+  iat: number; // Issued at (Unix timestamp)
+  exp: number; // Expiration (Unix timestamp)
+}
+
+/**
+ * Decode a JWT token without verifying the signature
+ */
+export function decodeToken(token: string): TokenPayload | null {
+  try {
+    const parts = token.split(".");
+    if (parts.length !== 3) {
+      return null;
+    }
+
+    const payload = parts[1];
+    const decoded = Buffer.from(payload, "base64url").toString("utf-8");
+    return JSON.parse(decoded) as TokenPayload;
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Get the expiration date from a token
+ */
+export function getTokenExpiry(token: string): Date | null {
+  const payload = decodeToken(token);
+  if (!payload || !payload.exp) {
+    return null;
+  }
+
+  return new Date(payload.exp * 1000);
+}
+
+/**
+ * Check if a token is expired
+ */
+export function isTokenExpired(token: string): boolean {
+  const expiry = getTokenExpiry(token);
+  if (!expiry) {
+    return true;
+  }
+
+  return expiry.getTime() <= Date.now();
+}
+
+/**
+ * Check if a token will expire within the given milliseconds
+ */
+export function isTokenExpiringSoon(token: string, withinMs: number): boolean {
+  const expiry = getTokenExpiry(token);
+  if (!expiry) {
+    return true;
+  }
+
+  return expiry.getTime() <= Date.now() + withinMs;
+}
+
+/**
+ * Get the agent ID from a token
+ */
+export function getAgentIdFromToken(token: string): string | null {
+  const payload = decodeToken(token);
+  return payload?.sub ?? null;
+}
+
+/**
+ * Get time remaining until token expires in milliseconds
+ */
+export function getTokenTimeRemaining(token: string): number {
+  const expiry = getTokenExpiry(token);
+  if (!expiry) {
+    return 0;
+  }
+
+  return Math.max(0, expiry.getTime() - Date.now());
+}

package/vitest.config.ts

@@ -0,0 +1,12 @@
+import { defineConfig } from "vitest/config";
+
+export default defineConfig({
+  test: {
+    globals: true,
+    environment: "node",
+    include: ["src/__tests__/**/*.test.ts"],
+    coverage: {
+      reporter: ["text", "json", "html"],
+    },
+  },
+});