refactor: a2a

backend/app/models/a2a.py

@@ -1,9 +1,34 @@
-from sqlalchemy import Boolean, Column, DateTime, ForeignKey, Integer, String, Text, func
+from sqlalchemy import Boolean, Column, DateTime, ForeignKey, Integer, String, Text, JSON, Enum as SQLEnum, func
 from sqlalchemy.orm import relationship
+import enum
 
 from app.database import Base
 
 
+class A2ATaskState(str, enum.Enum):
+    SUBMITTED = "SUBMITTED"
+    WORKING = "WORKING"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    CANCELED = "CANCELED"
+    INPUT_REQUIRED = "INPUT_REQUIRED"
+    AUTH_REQUIRED = "AUTH_REQUIRED"
+    REJECTED = "REJECTED"
+
+
+class A2APartType(str, enum.Enum):
+    TEXT = "text"
+    RAW = "raw"
+    URL = "url"
+    DATA = "data"
+
+
+class A2AMessageRole(str, enum.Enum):
+    USER = "user"
+    AGENT = "agent"
+    SYSTEM = "system"
+
+
 class A2ARemoteAgent(Base):
     __tablename__ = "a2a_remote_agents"
 
@@ -13,6 +38,17 @@ class A2ARemoteAgent(Base):
     base_url = Column(String, nullable=False)
     auth_scheme = Column(String, nullable=False, default="none")
     auth_token = Column(String, nullable=True)
+    shared_secret = Column(String, nullable=True)
+    mtls_ca_cert = Column(Text, nullable=True)
+    mtls_client_cert = Column(Text, nullable=True)
+    mtls_client_key = Column(Text, nullable=True)
+    oauth2_client_id = Column(String, nullable=True)
+    oauth2_client_secret = Column(String, nullable=True)
+    oauth2_token_url = Column(String, nullable=True)
+    oauth2_scopes = Column(String, nullable=True)
+    oidc_issuer_url = Column(String, nullable=True)
+    oidc_client_id = Column(String, nullable=True)
+    oidc_client_secret = Column(String, nullable=True)
     protocol_version = Column(String, nullable=True)
     capabilities_json = Column(Text, nullable=False, default="[]")
     card_json = Column(Text, nullable=True)
@@ -27,27 +63,84 @@ class A2ARemoteAgent(Base):
     project = relationship("Project")
 
 
+class A2APart(Base):
+    __tablename__ = "a2a_parts"
+
+    id = Column(Integer, primary_key=True, index=True)
+    message_id = Column(Integer, ForeignKey("a2a_messages.id", ondelete="CASCADE"), nullable=True, index=True)
+    artifact_id = Column(Integer, ForeignKey("a2a_artifacts.id", ondelete="CASCADE"), nullable=True, index=True)
+    part_type = Column(SQLEnum(A2APartType), nullable=False)
+    text_content = Column(Text, nullable=True)
+    raw_content = Column(Text, nullable=True)
+    url_content = Column(String, nullable=True)
+    data_content = Column(Text, nullable=True)
+    media_type = Column(String, nullable=True)
+    filename = Column(String, nullable=True)
+    metadata_json = Column(Text, nullable=False, default="{}")
+    created_at = Column(DateTime, default=func.now())
+
+    message = relationship("A2AMessage", back_populates="parts", foreign_keys=[message_id])
+    artifact = relationship("A2AArtifact", back_populates="parts", foreign_keys=[artifact_id])
+
+
+class A2AMessage(Base):
+    __tablename__ = "a2a_messages"
+
+    id = Column(Integer, primary_key=True, index=True)
+    message_id = Column(String, nullable=False, unique=True, index=True)
+    context_id = Column(String, nullable=True, index=True)
+    task_id = Column(String, ForeignKey("a2a_tasks.id", ondelete="CASCADE"), nullable=True, index=True)
+    role = Column(SQLEnum(A2AMessageRole), nullable=False)
+    extensions_json = Column(Text, nullable=False, default="{}")
+    reference_task_ids_json = Column(Text, nullable=False, default="[]")
+    created_at = Column(DateTime, default=func.now(), index=True)
+
+    task = relationship("A2ATask", back_populates="messages", foreign_keys=[task_id])
+    parts = relationship("A2APart", back_populates="message", cascade="all, delete-orphan")
+
+
+class A2AArtifact(Base):
+    __tablename__ = "a2a_artifacts"
+
+    id = Column(Integer, primary_key=True, index=True)
+    artifact_id = Column(String, nullable=False, unique=True, index=True)
+    task_id = Column(String, ForeignKey("a2a_tasks.id", ondelete="CASCADE"), nullable=False, index=True)
+    name = Column(String, nullable=True)
+    description = Column(Text, nullable=True)
+    metadata_json = Column(Text, nullable=False, default="{}")
+    extensions_json = Column(Text, nullable=False, default="{}")
+    created_at = Column(DateTime, default=func.now())
+    updated_at = Column(DateTime, default=func.now(), onupdate=func.now())
+
+    task = relationship("A2ATask", back_populates="artifacts")
+    parts = relationship("A2APart", back_populates="artifact", cascade="all, delete-orphan")
+
+
 class A2ATask(Base):
     __tablename__ = "a2a_tasks"
 
     id = Column(String, primary_key=True, index=True)
     project_id = Column(Integer, ForeignKey("projects.id"), nullable=False, index=True)
     tenant_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
+    context_id = Column(String, nullable=True, index=True)
     source = Column(String, nullable=False, default="local")
     remote_agent_id = Column(Integer, ForeignKey("a2a_remote_agents.id"), nullable=True, index=True)
     idempotency_key = Column(String, nullable=True, index=True)
-    state = Column(String, nullable=False, index=True, default="SUBMITTED")
+    state = Column(SQLEnum(A2ATaskState), nullable=False, index=True, default=A2ATaskState.SUBMITTED)
     input_text = Column(Text, nullable=False, default="")
     output_text = Column(Text, nullable=True)
     error_message = Column(Text, nullable=True)
     compatibility_mode = Column(Boolean, nullable=False, default=True)
     metadata_json = Column(Text, nullable=False, default="{}")
+    history_length = Column(Integer, nullable=False, default=0)
     created_at = Column(DateTime, default=func.now())
     updated_at = Column(DateTime, default=func.now(), onupdate=func.now())
     finished_at = Column(DateTime, nullable=True)
 
     project = relationship("Project")
     remote_agent = relationship("A2ARemoteAgent")
+    messages = relationship("A2AMessage", back_populates="task", cascade="all, delete-orphan", foreign_keys=[A2AMessage.task_id])
+    artifacts = relationship("A2AArtifact", back_populates="task", cascade="all, delete-orphan")
 
 
 class A2ATaskEvent(Base):